[論文レビュー] DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

Victor Sanh, Lysandre Debut+2 more|arXiv (Cornell University)|Oct 2, 2019

Topic Modeling参考文献 18被引用数 4,569

ひとこと要約

この論文では、事前学習段階における知識蒸留を用いて、言語理解能力の97％を保持するが、より小型で高速かつ効率的なBERTのバージョン、DistilBERTを紹介している。マスクされた言語モデル化、蒸留、コサイン距離損失を組み合わせた三重損失関数を用いることで、モデルサイズを40％削減し、推論時間を60％短縮しながら、GLUEおよび下流タスクの各性能を高い水準で維持している。

ABSTRACT

As Transfer Learning from large-scale pre-trained models becomes more prevalent in Natural Language Processing (NLP), operating these large models in on-the-edge and/or under constrained computational training or inference budgets remains challenging. In this work, we propose a method to pre-train a smaller general-purpose language representation model, called DistilBERT, which can then be fine-tuned with good performances on a wide range of tasks like its larger counterparts. While most prior work investigated the use of distillation for building task-specific models, we leverage knowledge distillation during the pre-training phase and show that it is possible to reduce the size of a BERT model by 40%, while retaining 97% of its language understanding capabilities and being 60% faster. To leverage the inductive biases learned by larger models during pre-training, we introduce a triple loss combining language modeling, distillation and cosine-distance losses. Our smaller, faster and lighter model is cheaper to pre-train and we demonstrate its capabilities for on-device computations in a proof-of-concept experiment and a comparative on-device study.

研究の動機と目的

• 自然言語処理タスクで高い性能を維持しつつ、より小型で高速で、効率的な汎用言語モデルを開発すること。
• BERTのような大規模な事前学習モデルをトレーニングおよびデプロイする際の計算コストと環境への影響を軽減すること。
• 微調整やタスク固有の適応の段階でのみではなく、事前学習段階においても知識蒸留を検討すること。
• モバイル機器やエッジデバイスで効率的に推論を実行できるほどに小型化されたモデルを実現すること。
• 事前学習段階での蒸留が、強力な一般化能力と転移能力を備えたコンactなモデルを生成できることを示すこと。

提案手法

• DistilBERTは、BERTの蒸留版であり、層の数を半分に減らし、トークンタイプの埋め込みとプーラーを削除することで、40％のパラメータ削減を達成している。
• 学生モデルは、教師のBERT-baseモデルの層を交互に取り出して初期化することで、知識を保持し、収束を改善している。
• 三重損失関数は、マスクされた言語モデル化（MLM）損失、教師のソフトラベルからの蒸留損失、および隠れ状態の方向を一致させるためのコサイン埋め込み損失の3つの成分を組み合わせている。
• 蒸留損失は、温度調整されたソフトマックスを用いて、教師のソフトラベル分布から学生への知識の転送を実現している。
• トレーニングは、BERTと同一のWikipediaおよびBookCorpusデータセットを用い、大規模バッチの勾配蓄積と動的マスキングを実施している。
• トレーニングは8台の16GB V100 GPUを用い、約90時間実行され、フルBERTトレーニングと比較して大幅に計算予算を削減している。

実験結果

リサーチクエスチョン

• RQ1 事前学習段階での知識蒸留は、BERTの性能の大部分を保持する小型言語モデルを生成できるか？
• RQ2 高性能な蒸留モデルを達成するために、損失関数のどの要素が不可欠か？
• RQ3 タスク固有の蒸留と比較して、事前学習段階での蒸留は下流タスクの性能にどのように影響するか？
• RQ4 蒸留モデルは、モバイル機器やエッジデバイスで効率的に実行できるほどに小型かつ高速にできるか？
• RQ5 教師の重みから学生を初期化することで、収束性と性能が向上するか？

主な発見

• DistilBERTは、GLUEベンチマークのマクロスコアでBERT-baseの97％の性能を達成しており、パラメータは40％にまで削減されている。
• STS-Bタスクでは、DistilBERTの推論速度はBERT-baseの60％速く、CPU上での推論時間は410msで、BERTの668msを上回っている。
• IMDbのセンチメント分類タスクでは、DistilBERTは92.82％の精度を達成しており、BERT-baseの93.46％からわずか0.64％の差にとどまっている。
• SQuAD 1.1では、DistilBERTはEM/F1で77.7/85.8を達成しており、BERT-baseの81.2/88.5から3.5ポイント以内の差にとどまっている。
• アブレーションスタディの結果、三重損失のいずれの成分（MLM、蒸留、コサイン損失）を削除しても性能が低下し、特にコサイン損失が顕著な貢献を示している。
• プロトタイプモバイルアプリケーションでは、iPhone 7 PlusでDistilBERTはBERT-baseよりも71％速く実行され、モデルサイズは207MBであった。