[論文レビュー] AutoDistill: an End-to-End Framework to Explore and Distill Hardware-Efficient Language Models

Xiaofan Zhang, Zongwei Zhou+2 more|arXiv (Cornell University)|Jan 21, 2022

Topic Modeling被引用数 5

ひとこと要約

AutoDistill は、ベイジアン最適化を用いた多目的ニューラルアーキテクチャサーチと、迅速なモデル候補評価のためのフラッシュディスティルレーションを組み合わせたエンドツーエンドフレームワークであり、ハードウェアに最適化されたタスクに依存しない NLP モデルを構築する。このフレームワークは、MobileBERT よりも最大 3.2% の高い事前学習精度と 1.44 倍の高速な推論遅延を達成しており、20.6M パラメータのモデルが GLUE および SQuAD ベンチマークで BERT BASE や他のディスティル化モデルを上回っている。

ABSTRACT

Recently, large pre-trained models have significantly improved the performance of various Natural LanguageProcessing (NLP) tasks but they are expensive to serve due to long serving latency and large memory usage. To compress these models, knowledge distillation has attracted an increasing amount of interest as one of the most effective methods for model compression. However, existing distillation methods have not yet addressed the unique challenges of model serving in datacenters, such as handling fast evolving models, considering serving performance, and optimizing for multiple objectives. To solve these problems, we propose AutoDistill, an end-to-end model distillation framework integrating model architecture exploration and multi-objective optimization for building hardware-efficient NLP pre-trained models. We use Bayesian Optimization to conduct multi-objective Neural Architecture Search for selecting student model architectures. The proposed search comprehensively considers both prediction accuracy and serving latency on target hardware. The experiments on TPUv4i show the finding of seven model architectures with better pre-trained accuracy (up to 3.2% higher) and lower inference latency (up to 1.44x faster) than MobileBERT. By running downstream NLP tasks in the GLUE benchmark, the model distilled for pre-training by AutoDistill with 28.5M parameters achieves an 81.69 average score, which is higher than BERT_BASE, DistillBERT, TinyBERT, NAS-BERT, and MobileBERT. The most compact model found by AutoDistill contains only 20.6M parameters but still outperform BERT_BASE(109M), DistillBERT(67M), TinyBERT(67M), and MobileBERT(25.3M) regarding the average GLUE score. By evaluating on SQuAD, a model found by AutoDistill achieves an 88.4% F1 score with 22.8M parameters, which reduces parameters by more than 62% while maintaining higher accuracy than DistillBERT, TinyBERT, and NAS-BERT.

研究の動機と目的

• データセンターサービスにおける自動的でスケーラブルかつ遅延を考慮したモデルディスティルレーションの欠如に対処すること。
• 大規模事前学習モデルのための学生モデルアーキテクチャの探索にかかるコストと時間を削減すること。
• 完全なディスティルレーションの前段階で、有望な学生モデルを効率的に同定できること。
• ターゲットハードウェア上で予測精度と推論遅延を同時に最適化すること。
• データセンターサービスに適したエンドツーエンドで自動的かつ生産環境対応の可能な NLP モデルのディスティルレーションを支援すること。

提案手法

• ベイジアン最適化（BO）を用いて、ターゲットハードウェア上の精度と推論遅延の両方を同時に最適化する多目的ニューラルアーキテクチャサーチ（NAS）を実施する。
• フラッシュディスティルレーションを導入し、全ディスティルレーションステップの 5%（37k ステップ vs. 740k）のわずかなステップ数で、有望な学生モデル候補を早期に特定するモデルに依存しない手法を採用する。
• プログレッシブな知識伝達戦略をフラッシュディスティルレーションで採用し、プログレッシブと事前学習のディスティルレーションステップの比率を固定（0.48）に設定する。
• コスト効率の良い選定のため、対数的な精度上昇曲線を用いて候補モデルを早期に評価し、最適な停止ポイントを特定する。
• BO とフラッシュディスティルレーションを統合することで、品質の高いモデル発見を維持しながらも、探索プロセスを加速する。
• TPUv4i ハードウェア上でエンドツーエンドのディスティルレーションを実施し、GLUE および SQuAD ベンチマークで包括的な性能評価を実施する。

実験結果

リサーチクエスチョン

• RQ1 自動的でエンドツーエンドのフレームワークは、最小限の人的介入でハードウェアに最適化された NLP モデルを効果的に探索・ディスティル化できるか？
• RQ2 多目的ニューラルアーキテクチャサーチは、モデルディスティルレーションにおける精度と推論遅延のバランスをどのように効率的に適用できるか？
• RQ3 短時間のプログレッシブディスティルレーションフェーズ（フラッシュディスティルレーション）は、完全なトレーニングの前段階で有望な学生モデルを信頼性高く特定できるか？
• RQ4 フラッシュディスティルレーションは、最終的なモデル性能を損なわせることなく、モデル探索のコストをどの程度削減できるか？
• RQ5 得られたディスティル化モデルは、標準ベンチマークにおいて、既存の最先端のディスティル化モデルを精度と遅延の両面で上回ることができるか？

主な発見

• AutoDistill は、TPUv4i 上で MobileBERT より最大 3.2% の高い事前学習精度と最大 1.44 倍の高速な推論遅延を達成した 7 つの学生モデルアーキテクチャを発見した。
• 28.5M パラメータのモデルは 81.69 の平均 GLUE スコアを達成し、BERT BASE（109M）、DistillBERT（67M）、TinyBERT（67M）、NAS-BERT、MobileBERT（25.3M）を上回った。
• 最もコンパクトなモデル（20.6M パラメータ）は、サイズが 81.1% 減少したにもかかわらず、BERT BASE、DistillBERT、TinyBERT、MobileBERT を上回る平均 GLUE スコアを達成した。
• SQuAD ベンチマークでは、23M パラメータ未満の 2 つの AutoDistill モデルが、DistillBERT、TinyBERT、NAS-BERT を上回る F1 スコアを達成した一方で、パラメータ数を 62% 以上削減した。
• 37k ステップのフラッシュディスティルレーションは、トレーニングコストの 5% で有望なモデルを効果的に特定でき、全ディスティルレーションと比較して安定した相対的な性能順位を維持した。
• BO とフラッシュディスティルレーションの統合により、全体の探索コストを顕著に削減した一方で、精度と遅延の両面で既存のベースラインを上回るモデルの発見が可能になった。