[論文レビュー] Rethinking Network Pruning -- under the Pre-train and Fine-tune Paradigm

Dongkuan Xu, Ian En-Hsu Yen+2 more|arXiv (Cornell University)|Apr 18, 2021

Topic Modeling参考文献 27被引用数 6

ひとこと要約

この論文では、事前学習済みのBERTを教師モデルとして用い、微調整段階でスパースなプルーニングを適用する知識に配慮したスパースプルーニング手法、SparseBERTを提案する。知識蒸留を用いて一般知識とタスク固有の知識を両方保持する。GLUEベンチマークで平均1.4%の精度低下で20倍のパラメータおよびFLOPs圧縮を達成し、構造的プルーニングや蒸留ベースのモデルを含む先行手法を上回る性能を示した。

ABSTRACT

Transformer-based pre-trained language models have significantly improved the performance of various natural language processing (NLP) tasks in the recent years. While effective and prevalent, these models are usually prohibitively large for resource-limited deployment scenarios. A thread of research has thus been working on applying network pruning techniques under the pretrain-then-finetune paradigm widely adopted in NLP. However, the existing pruning results on benchmark transformers, such as BERT, are not as remarkable as the pruning results in the literature of convolutional neural networks (CNNs). In particular, common wisdom in pruning CNN states that sparse pruning technique compresses a model more than that obtained by reducing number of channels and layers (Elsen et al., 2020; Zhu and Gupta, 2017), while existing works on sparse pruning of BERT yields inferior results than its small-dense counterparts such as TinyBERT (Jiao et al., 2020). In this work, we aim to fill this gap by studying how knowledge are transferred and lost during the pre-train, fine-tune, and pruning process, and proposing a knowledge-aware sparse pruning process that achieves significantly superior results than existing literature. We show for the first time that sparse pruning compresses a BERT model significantly more than reducing its number of channels and layers. Experiments on multiple data sets of GLUE benchmark show that our method outperforms the leading competitors with a 20-times weight/FLOPs compression and neglectable loss in prediction accuracy.

研究の動機と目的

• 事前学習済み言語モデル（例：BERT）におけるスパースプルーニングの効果性に課題があり、スパースプルーニングは構造的プルーニングに比べて性能が劣ることがあることに対処する。
• プルーニング中に事前学習による汎用的言語知識と微調整によるタスク固有の知識を両方保持する。
• スパースプルーニングが、チャネル/レイヤー削減と比較して、BERTにおいて顕著に高い圧縮比を達成できることを示す。
• 特にリソース制限のある環境においても、極めて高いスパース性を維持しながら高い性能を発揮するプルーニング戦略を開発する。
• 標準ベンチマーク（GLUE、SQuAD）およびスパース推論をネイティブにサポートするハードウェアプラットフォーム上で、手法の有効性を検証する。

提案手法

• プルーニングは微調整段階で実施され、事前学習済みBERTを学生モデルとして用い、微調整済みBERTを教師モデルとして用いる。
• 知識蒸留を用いて、ログィット蒸留と損失最小化を通じて、教師モデルからプルーニングされた学生モデルへタスク固有の知識を転送する。
• スパース性は自己注意およびフィードフォワード層の線形変換に適用され、これらは過パrameter化されており、計算負荷が大きいと特定された。
• 構造的・表現的整合性を保つために、注意ヘッドや埋め込み層のプルーニングを回避する。
• プルーニングプロセスは微調整ループに統合されており、スパース性と性能の共同最適化を可能にする。
• Moffett AIのANTOMプラットフォーム（ネイティブなスパーステンソルサポートを備える）を用いて、実世界の推論高速化を測定するハードウェア評価を実施する。

実験結果

リサーチクエスチョン

• RQ1 事前学習と微調整のパラダイム下で、スパースプルーニングがチャネル削減などの構造的プルーニングよりも高い圧縮比を達成できるか？
• RQ2 事前学習と微調整からの両方の知識をプルーニング中に保持するにはどうすればよいか？
• RQ3 微調整段階での知識蒸留により、極めて高いスパース性でも効果的なスパースプルーニングが可能になり、高い精度を維持できるか？
• RQ4 事前学習段階でのプルーニングや蒸留なしの微調整段階でのプルーニングと比較して、スパースプルーニングの性能はどのように異なるか？
• RQ5 ネイティブなスパーステンソルサポートを備えたプラットフォーム上でのスパースBERTモデルの実際のハードウェア推論高速化はどの程度か？

主な発見

• SparseBERTは、BERT-baseで20倍の圧縮（残存パラメータ5%）を達成し、GLUE開発セット全体で平均1.4%の精度低下にとどまる。
• MRPCデータセットでは、パラメータを85.53Mから4.84M、FLOPsを10.87 GFLOPsから0.54 GFLOPsに削減し、20倍の圧縮を達成した。
• GLUEベンチマークにおいて、TinyBERT、DistilBERT、BERT-of-Theseusを含むすべての主要な競合手法を、同等または低いスパース性で上回った。
• SQuAD v1.1およびv2.0において、下流の微調整段階でのプルーニングよりも、特に高スパース性条件下で一貫して優れた性能を示した。
• Moffett AIのANTOMハードウェア上では、20倍の圧縮で10倍以上の推論高速化を達成し、4倍圧縮までに近似線形の高速化が得られた。
• 同等のスパース性において、事前学習段階でのプルーニングに比べ、SparseBERTはより優れた適合能力を示した。これは、RTEタスクで同等のスパース性でも訓練損失が低く、開発セット精度が高いことで裏付けられた。