Block Pruning For Faster Transformers

https://arxiv.org/abs/2109.04838
EMNLP 2021

事前学習を行うことで、分類と生成の両方のタスクにおいてモデルの精度が向上しますが、その代償として、モデルのサイズが大きくなり、処理速度も低下します。刈り込み手法は、モデルサイズを小さくする効果的な方法であることが証明されており、一方、蒸留手法は推論を高速化する方法であることが証明されている。本論文では、小さくて高速なモデルの両方を対象とした、ブロック・プルーニング手法を紹介する。このアプローチは、任意のサイズのブロックを考慮することで構造化手法を拡張し、この構造を微調整のための運動刈り込みパラダイムに統合する。このアプローチは、注意ヘッドなどの基礎的なモデルの完全なコンポーネントを刈り取ることを学習することを発見しました。実験では、分類と生成のタスクを検討し、SQuAD v1のBERTを2.4倍高速化、74%縮小し、F1を1%低下させた刈り込みモデルが得られました。

https://github.com/huggingface/nn_pruning

1 はじめに

訓練済みの変換モデルは、分類タスクと生成タスクの両方において、NLPタスクの標準となっています(Devlin et al., 2019; Lewis et al., 2020)。最近の傾向として、標準的なベンチマークでの性能向上をもたらしながらも、モデルのサイズは大きくなり続けています(Rosset, 2020)。このような状況から、学習済みモデルの保存サイズを縮小し、効率を高める必要性が高まっています。特定のタスクのために微調整されたモデルのストレージサイズを削減するには、プルーニング手法が非常に有効であることがわかっています。マグニチュード・プルーニング（Han et al., 2015）、L0正則化（Louizos et al., 2018）、ロッタリー・チケット仮説（Frankle and Carbin, 2018）、diffプルーニング（Guo et al., 2020）、ムーブメント・プルーニング（Sanh et al., 2020）などのアプローチは、モデルサイズの顕著な削減を実証しています。 Movement pruningは、SQuAD v1.1で1%の精度低下に対して77%のパラメータストレージの節約を実現しています。しかし、これらのモデルは、標準的なハードウェアで実行するためには、しばしば元の密な形状を再構築する必要があるため、実際の効率化はほとんど得られません。一方、蒸留法は、DistilBERT（Sanh et al.、2019年）、TinyBERT（Jiao et al.、2019年）、またはMobileBERT（Sun et al.、2020年）によって示されているように、より高速なモデルを生成するのに効果的である。これらのアプローチは、対象となる蒸留を利用して、標準的なハードウェアで高速に動作する緻密な構造を持つ小さなモデルを生成します。しかし、慎重なエンジニアリングとサイズ選択を行わないと、これらのモデルは刈り込まれたものよりもはるかに大きくなります。

本研究では、ブロック・プルーニングによってこのギャップを埋めることを目標としています。このアプローチは、個々のパラメータの刈り込みとは異なり、高密度なハードウェアで最適化できる刈り込みを推奨しています。これは、構造化されたアプローチ（McCarley, 2019）で一般的に使用される行または列ベースのプルーニングよりも剛性が低いアプローチであり、変圧器に効果的に適用することは困難でした。このアプローチを、微調整中に事前に訓練されたモデルをプルーニングするためのシンプルな手法であるMovement pruning（Sanh et al.、2020）と統合する。最終的な手法1では、追加のハイパーパラメータやトレーニング要件はほとんどない。

実験では、さまざまなベンチマークデータセットを用いて、精度と効率を比較しました。その結果、学習時にsub rowの正方形ブロックを利用しているにもかかわらず、この手法はモデルの完全なコンポーネントを排除するように学習し、多数のアテンションヘッドを効果的に削除するという驚くべき結果が得られました。この効果により、標準的な構造化されたフィードフォワード層の刈り込みよりも高速化することができました。その結果、SQuAD v1.1ではF1を1%低下させながら2.4倍の高速化を、QQPではF1を1%低下させながら2.3倍の高速化を達成しました。また、CNN/DailyMailにおける要約の実験では、ROUGEの全指標で平均2ポイントの低下で1.39倍の高速化を達成し、デコーダの重みを3.5倍に減らすことができました。

3 背景

パラメーターθを持つトランスフォーマーモデルから始めて、我々の目標は、特定のエンドタスクのために微調整され、かつ並列ハードウェア上で効率的に推論が計算できるように小さくなったパラメーターθ0のセットを作り出すことです。トランスフォーマーのパラメータバジェットの中で最も大きなラインは、フィードフォワードネットワークサブレイヤー(FFN)とマルチヘッドアテンションサブレイヤー(MHA)の2つです。 FFNのパラメータは、R dmodel×dffとR dff×dmodelという転置された形の2つの行列（W1とW2）で構成され、dmodelは隠れたサイズ、dff dmodelは内側のサイズです。

これらは、ネットワークが標準的に使用するものです。 MHAのパラメータは、R dmodel×dmodel（query, key, value, out）のサイズの4つの投影行列（Wq, Wk, Wv, Wo）で構成されています。これらは，隠れたベクトルを注目部分のコンポーネントとの間で投影するために使用されます．実装では、この投影は、折りたたまれたテンソル形式の行列 R nheads×(dmodel /nheads) ×dmodel で行われます。標準的な微調整では、θから始めて、損失L（例えば、分類のためのクロス・エントロピー）を最適化します。

スコアベースのプルーニング法（Ramanujan et al, 2019）は、各パラメータiに対してスコアパラメータSを導入することでモデルを修正し、元のパラメータマトリクスをマスクされたバージョンW' = W elementwise_prod M(S)に置き換えます。例えば、マグニチュード・プルーニングの最も単純なバージョンでは、マスクは絶対値が低いパラメータをゼロアウトするだけである。 Movement pruning (Sanh et al., 2020)は、モデルがこれらのスコアパラメータを最適化することを促すスコアベースのプルーニングアプローチである。具体的には、閾値パラメータτに対してM(S) = 1(S > τ )を設定し、正則化された目的を最適化する動きプルーニングのソフトムーブメントバリアントに注目する。

この刈り込み目的は，モデルがパラメータを微調整することを促す一方で，重要でないパラメータのスコアを下げることで，より多くのスパース性を促す．閾値を超えて学習するためには，ストレートスルー推定法（Bengio et al. が使用される

Movement pruningとdistillationの組み合わせは、既存のモデルのパラメータ数を減らすのに非常に効果的な方法であることが示されており、SQuAD v1.1のF1が87.5の場合、テストでは94%のpruningが得られました（BERT-baseは88.5）。この結果、蒸留法のみの場合と比べて、モデルが大幅に小さくなりました。しかし、このようなスパース度であっても、このスタイルのスパースな行列-ベクトル積を十分に活用できないほとんどの標準的なハードウェアで実行した場合、モデルは大幅に高速化されません。

4 モデル :ブロック移動プルーニング

本研究では，移動プルーニングを拡張して，ローカルパラメータのブロックで動作するようにします．具体的には、変換器の各行列を固定サイズのブロックに分割します。この設定は、非構造化手法の恣意的なプルーニングを超えたもので、データの局所性を効率化に必要なものに近づけることを目的としています2。各パラメータ行列W∈RM×Nに対して、固定サイズのブロック構造（M' , N' ）を仮定します。これらのブロックは、対応するスコア行列S∈RM/M'×N/N'から得られる共通のスコアパラメータを持つ正則化の個々のグループとして動作します。マスキングされた重みの計算は、閾値を拡大することで行われ、すなわち

過去の研究と同様に，このモデルは教師モデルの性能と一致するように distillation で学習されます．新しいモデル構造を完全に指定する必要がある他の蒸留法とは異なり、本手法ではブロックのサイズと形状、すなわちモデルの各パラメータ行列に対する（M' , N' ）のセットのみを必要とします。ブロックが大きすぎると剪定が難しくなりますが、小さすぎると効率的な推論ができません。

探索空間を小さくするために、(M' , N' ) att と (M' , N' ) ff に試すものを限定する。一方の注意ウェイト Wq, Wk, Wv, Wo と、他方のフィードフォワードウェイト W1, W2 については、すべての層で同じブロックサイズを使用することにする。動きの刈り込み正則化項を次のように分割します。

これにより、スコアパラメータが受け取るグラデーションの違いを考慮することができます。探索空間をさらに小さくするために、2種類のブロックでテストを行います。

(32, 32) : square blocks (Block)
(1, dmodel) and (dmodel, 1) : paired FFNの行と列(Dim)でのdimension pruning

これらのブロックサイズは、効率的なモデルを可能にします。少なくとも(16, 16)サイズのブロックは、適切なGPUカーネルで効率的に計算することができますが、行、列、または頭部を行列から完全に取り除くことができ、残った行列は密なものになります。このアプローチを検証するために、ベースラインとなる2つのブロックタイプを追加しました。

(2^n, 2^n), n in [2,5] : 2つのブロックサイズの小さい方のべき乗．性能におけるサイズの影響を調べるため(Block)
(d_model/ n_heads, d_model): attention heads用(Heads)

最初のものは小さいblockを考慮し，二番目は非常に大きいfunctional blockを考慮する

e4exp / paper_manager_abstract