[論文レビュー] Prune Once for All: Sparse Pre-Trained Language Models

Ofir Zafrir, Ariel Larey+3 more|arXiv (Cornell University)|Nov 10, 2021

Topic Modeling被引用数 26

ひとこと要約

この論文は Prune Once for All（Prune OFA）を紹介します。これは重み pruning と蒸留を統合して sparse な事前学習済み Transformer 言語モデルを訓練する、アーキテクチャに依存しない手法です。高い疎性（例: 85–90%）を、下流タスクでの精度低下を最小限に抑えつつ実現し、任意の量子化を追加することも可能です。

ABSTRACT

Transformer-based language models are applied to a wide range of applications in natural language processing. However, they are inefficient and difficult to deploy. In recent years, many compression algorithms have been proposed to increase the implementation efficiency of large Transformer-based models on target hardware. In this work we present a new method for training sparse pre-trained Transformer language models by integrating weight pruning and model distillation. These sparse pre-trained models can be used to transfer learning for a wide range of tasks while maintaining their sparsity pattern. We demonstrate our method with three known architectures to create sparse pre-trained BERT-Base, BERT-Large and DistilBERT. We show how the compressed sparse pre-trained models we trained transfer their knowledge to five different downstream natural language tasks with minimal accuracy loss. Moreover, we show how to further compress the sparse models' weights to 8bit precision using quantization-aware training. For example, with our sparse pre-trained BERT-Large fine-tuned on SQuADv1.1 and quantized to 8bit we achieve a compression ratio of $40$X for the encoder with less than $1\%$ accuracy loss. To the best of our knowledge, our results show the best compression-to-accuracy ratio for BERT-Base, BERT-Large, and DistilBERT.

研究の動機と目的

• モデルサイズの拡大と環境コストの高まりにより、効率的でデプロイ可能な Transformer LMs の必要性を動機づける。
• アーキテクチャに依存しない手法（Prune OFA）を提案し、転移学習機能を保持する sparse な事前学習済みLMを訓練する。
• 疎な事前学習モデルが、複数の下流タスクで最小限の精度低下でファインチューニング可能であることを示す。
• その後の量子化がさらなるモデルサイズ削減をもたらし、精度への影響を控えめに抑えることを示し、再現性のあるツールとモデルを公開する。

提案手法

• 単一の事前訓練/知識伝達プロセス中に非構造的な重み prune を用いて sparse な事前学習済み LM を得る。
• pruning 過程で Gradual Magnitude Pruning（GMP）を Learning Rate Rewinding（LRR）および知識蒸留（KD）と組み合わせる。
• downstream のファインチューニング時に疎性パターンを保持するための Pattern-Lock メカニズムを適用する。
• 英語版 Wikipedia で事前学習を行い、SQuADv1.1、GLUE タスクなどの下流タスクへ KD を用いて性能を保持した転送を行う。
• 任意で Quantization-Aware Training（QAT）を適用し、8-bit の量子化 sparse モデルを得る。
• 再現性のためのスクリプトと sparse な事前学習済みモデルを含むオープンソース圧縮ライブラリを提供する。

実験結果

リサーチクエスチョン

• RQ1 事前学習中の pruning によって、下流タスクへ転送可能な sparse な言語モデルを、最小限の精度低下で得られるか？
• RQ2 GMP、LRR、KD を pruning 中に組み合わせると、タスク固有 pruning より転送性能が向上するか？
• RQ3 パターンロックで疎性パターンを保持することは、ファインチューニング中の精度維持に役立つか？
• RQ4 下流の量子化（8-bit QAT）は、 sparse な事前学習モデルの圧縮と精度のトレードオフにどのように影響するか？

主な発見

モデル	疎性	KD を用いた転送	SQuAD	MNLI (m/mm)	SST-2	QNLI	QQP
Prune OFA (BERT-Base)	85%	Yes	78.59	86.63	81.67	82.53	91.34	89.95
Prune OFA (BERT-Base)	85%	No	78.00	86.16	82.45	83.05	88.82	87.79
Prune OFA (BERT-Base)	85%	Yes	81.10	88.42	82.71	83.67	91.46	90.34

• Prune OFA は高い疎性（85–90%）を達成しつつ、SQuADv1.1 および GLUE タスクでの転送性能が密なベースラインおよび従来の prune 手法と比較して競争力を維持する。
• 転送時に KD を使用すると結果が改善され、LRR と pattern-lock を KD と組み合わせるとさらなる改善が得られ、多くのタスクで精度の低下は最小限に抑えられる。
• sparse モデルに対する量子化対応訓練（QAT）は、 full-precision sparse モデルと比較して平均約 0.67% の相対的な精度低下という小さな影響で、モデルサイズを大幅に削減する圧縮効率を生み出す。
• BERT-Large を 90% の疎性で用いた場合、ほとんどのタスクで約 1% 前後の精度低下に収まり、パラメータ効率（非ゼロパラメータ数）の点では密な BERT-Base を上回る。
• 著者らは圧縮ライブラリと sparse な事前学習済みモデルを公開し、モデル pruning と圧縮における再現性のある研究を促進する。