こんにちは。 リサーチャーの勝又です。 私はレトリバで自然言語処理、とくに要約や文法誤り訂正に関する研究の最新動向の調査・キャッチアップなどを行っております。
ニューラルネットワークモデルの軽量化や推論高速化手法として、蒸留を利用した小さいモデル作成が挙げられます。 今回はtask-specific BERTの蒸留をDeepSpeed Compressionで試してみようと思います。
蒸留を用いたBERTの軽量化とは
近年、自然言語処理の分野ではBERT（Bidirectional Encoder Representations from Transformers）1を用いたモデルが盛んに研究されています。 BERTは確かに精度の面では優れているのですが、その分モデルが大きく、推論速度に難がある事が知られています。
知識蒸留（Knowledge Distillation2; 本記事では蒸留と呼称します）は、簡単に説明すると学習済みモデル（Teacherモデル）を用意して、そのモデルの出力を真似するように別のモデル（Studentモデル）を学習する手法を指します。 蒸留を用いてモデルサイズ軽減を行う場合は、Teacherモデルを大規模モデルに、Studentモデルを小規模モデルに設定する事が多いです。 蒸留の学習を行う際にはTeacherモデルの出力logitをStudentモデルのobjectiveに使用します。 また、BERTの蒸留の場合、出力logit以外にもhidden stateやattention matricesを蒸留に使用する事が有効である事が報告されています3。 詳細につきましてはぜひ元論文の方をご確認ください。
BERTの蒸留では、StudentモデルとしてTeacherモデルよりも層の数が少ないモデルを用意します。
英語モデルにはなりますが、Huggingface Model Hubに登録されているdistilbert-base-uncased4は、BERTのbaseモデル（12層）をTeacherとして作成した、6層のStudentモデルです。
このモデルはいわゆるtask-agnostic BERTの蒸留で、Masked Language Modelなどで事前学習されたBERTを元に作成したものです。
この作成したStudentモデルに対して解きたい下流タスクデータでfine-tuningする事で、12層のBERTに迫る精度を達成しつつ高速に推論する事ができます。
DeepSpeed Compressionを用いた蒸留
弊社技術ブログでも以前紹介したDeepSpeedに、v0.7.0からDeepSpeed Compressionが追加されました。 今回はこちらを使用してBERTの蒸留を行いたいと思います。
DeepSpeedとは
最初に、DeepSpeedについて簡単に説明します。 DeepSpeedの公式サイトには次の文言が記述されています。
DeepSpeed is an easy-to-use deep learning optimization software suite that enables unprecedented scale and speed for Deep Learning Training and Inference.
Deep Learningの学習や推論の高速化を行うライブラリで、2022年11月現在、DeepSpeedには次の3つのトピックが存在しています。
- Training
- Inference
- Compression
このDeepSpeed CompressionにはXTC5およびZeroQuant6と呼ばれるモデル軽量化の実装が行われています。
XTCの一部（元論文中の5章 Step I: Lightweight layer reduction.）として、蒸留を用いてBERTの層を減らす処理が実装されているので、今回はそちらを試して見たいと思います。
DeepSpeed Compressionを使って蒸留してみる
DeepSpeed Compressionの公式ExampleはこちらのGitHubに置かれています。 今回はこのExampleに則ってfine-tune済みのtask-specific BERTの蒸留を行いたいと思います。
task-specific BERTの用意
今回は先に12層のBERTを下流タスクでfine-tuneしておいて、そのtask-specificモデルをTeacherとして蒸留を行いたいと思います。 具体的にはLivedoorニュースコーパスを利用して、記事が与えられた時にその記事が属するカテゴリを分類するタスク（9クラス分類）でfine-tuningしました。 蒸留の際の学習データとしてもこのデータは使用するので消さないようにしてください。
DeepSpeedのConfigの作成
次に、DeepSpeedのConfigファイルを用意します。 DeepSpeed CompressionのConfigファイルのテンプレートはこちらにあるので適宜埋めると蒸留を行う事が可能になります。
今回は下記のようなConfigファイルを使用しました（少し長いので折りたたんでいます）。
今回の蒸留で重要なConfigは次の箇所です。
"layer_reduction": { "enabled": true, "keep_number_layer": 5, "module_name_prefix": "bert.encoder.layer", "teacher_layer": [ 2, 4, 6, 8, 10 ], "other_module_name": [ "bert.pooler", "bert.embeddings", "classifier" ] },
この部分ではStudentモデルの層を幾つにするか指定しています。
keep_number_layerでStudentモデルを5層にし、その5層の初期化にTeacherモデルの2, 4, 6, 8, 10層を使うことをteacher_layerで指定しています。
また、他にもこのConfigではDeepSpeedの学習に関する設定も行っています。 今回はDeepSpeed Compressionが中心なので、こちらの解説に関しては省略いたします。
蒸留スクリプトの実行
TeacherモデルとConfigファイルが用意できたら蒸留を行います。
こちらのスクリプトrun_glue_no_trainer.pyを動かすだけなのですが、実はこのスクリプトはあくまでGLUEタスクに対するものなので、自分で新規タスクを設定する場合は少し変更する必要があります。
こちらも詳細につきましては少し長くなるので折りたたんでいます。 ご興味のある方はご確認ください。
蒸留スクリプトの用意ができたら次のコマンドで蒸留を行います。
python -m torch.distributed.launch --nproc_per_node=1 \ run_glue_no_trainer.py \ --distill_method one_stage \ --model_name_or_path /path/to/fine-tuned-teacher \ --train_file /path/to/train.json \ --validation_file /path/to/valid.json \ --max_length 128 \ --pad_to_max_length \ --per_device_train_batch_size 32 \ --per_device_eval_batch_size 64 \ --num_train_epochs 3 \ --num_warmup_epochs 1 \ --eval_step 1000 \ --deepspeed_config /path/to/ds_config.json \ --deepspeed \ --save_best_model --clean_best_model \ --output_dir /path/to/save_dir
このスクリプトを実行すると、output_dirで指定したディレクトリにモデルなどが保存されます。
save_dir ├── clean │ ├── config.json │ ├── pytorch_model.bin │ └── vocab.txt └── ds_config.json
作成したStudentモデルはhuggingface/transformersで読み込む事が可能です。
from transformers import AutoModelForSequenceClassification model = AutoModelForSequenceClassification("/path/to/save_dir/clean")
蒸留実験
それでは、実際に蒸留を用いた実験を行おうと思います。
実験設定
前述の通り、今回はLivedoorニュースコーパスの記事に対するカテゴリを推定する分類タスクで検証を行いました。
Huggingface Model Hubのcl-tohoku/bert-base-japanese-whole-word-masking7に対して分類タスクをfine-tuningし、そのモデルをTeacher、前述のDeepSpeed Compressionを用いてStudentモデルを作成しました。
今回はこのTeacherモデルとStudentモデルを比較します。
なお、推論はCPU上で行い、batch sizeは1、推論時のmax lengthは512で行いました。
また推論速度8はテストデータそれぞれの平均を使用しています。
TeacherモデルとStudentモデルの比較
TeacherモデルとStudentモデルのモデルサイズ、推論速度、精度を計測した結果が次の表です。
| モデルサイズ（MB） | 推論速度（sec） | 精度（acc） | |
|---|---|---|---|
| Teacherモデル（12層） | 423 | 0.2374 | 0.96 |
| Studentモデル（5層） | 233 | 0.1088 | 0.93 |
モデルサイズは423MBから233MBと44.9%程度減少できました。 また、推論速度についても54.2%の速度減少ができました。 一方でモデルが小規模になったことにより精度も減少しています。 このように、日本語データでも英語の場合と同様に蒸留によるメリットデメリットが発生している事がわかります。
まとめ
今回の記事では、DeepSpeed Compressionを用いてBERTの蒸留を行いました。 簡単な検証ではありましたが、想定通りモデルサイズは減少し、推論は高速化し、精度が少し減少しました。
今回はrun_glue_no_trainer.pyを使って蒸留を行ったこともあり、task-specificになりましたが、少し手直しすればtask-agnostic BERTの蒸留も可能だと思います。
DeepSpeedの蒸留の良いところの1つとして、DeepSpeedの学習効率化の機能も使えるところがあると思うので、そちらをうまく使うと効率よくtask-agnostic BERTの蒸留ができそうです。
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova. [paper]↩
- Distilling the Knowledge in a Neural Network. Geoffrey Hinton, Oriol Vinyals, Jeff Dean. [paper]↩
- TinyBERT: Distilling BERT for Natural Language Understanding. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, Qun Liu. [paper]↩
- https://huggingface.co/distilbert-base-uncased↩
- Extreme Compression for Pre-trained Transformers Made Simple and Efficient. Xiaoxia Wu, Zhewei Yao, Minjia Zhang, Conglong Li, Yuxiong He. [paper]↩
- ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers. Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, Yuxiong He. [paper]↩
- https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking↩
- この推論速度にはモデルロードは含んでいません。テストデータをdata_loaderでモデルに入力し、すべてのテストデータが推論されるまでの時間を計測、1サンプルあたりの平均時間を計測しました。↩
