[論文レビュー] KroneckerBERT: Learning Kronecker Decomposition for Pre-trained Language Models via Knowledge Distillation

Marzieh S. Tahaei, Ella Charlaix+3 more|arXiv (Cornell University)|Sep 13, 2021

Topic Modeling参考文献 35被引用数 10

ひとこと要約

KroneckerBERT は、変換器ブロック内の埋め込み層およびすべての線形変換に対して、Kronecker分解と中間層知識蒸留を組み合わせることで、BERTベースのモデルの極端な圧縮を実現する新しいフレームワークを提案する。19倍の圧縮比（BERT_BASEの5％）において、GLUEおよびSQuADで最先端の性能を達成しており、高い圧縮率にもかかわらず、先行手法を上回っている。

ABSTRACT

The development of over-parameterized pre-trained language models has made a significant contribution toward the success of natural language processing. While over-parameterization of these models is the key to their generalization power, it makes them unsuitable for deployment on low-capacity devices. We push the limits of state-of-the-art Transformer-based pre-trained language model compression using Kronecker decomposition. We use this decomposition for compression of the embedding layer, all linear mappings in the multi-head attention, and the feed-forward network modules in the Transformer layer. We perform intermediate-layer knowledge distillation using the uncompressed model as the teacher to improve the performance of the compressed model. We present our KroneckerBERT, a compressed version of the BERT_BASE model obtained using this framework. We evaluate the performance of KroneckerBERT on well-known NLP benchmarks and show that for a high compression factor of 19 (5% of the size of the BERT_BASE model), our KroneckerBERT outperforms state-of-the-art compression methods on the GLUE. Our experiments indicate that the proposed model has promising out-of-distribution robustness and is superior to the state-of-the-art compression methods on SQuAD.

研究の動機と目的

• 低容量のエッジデバイスに大規模で過パrameter化された事前学習済み言語モデルをデプロイする課題に対処すること。
• 顕著な性能低下を伴わずに、10倍以上の圧縮比を実現する圧縮フレームワークの開発。
• Kronecker分解を標準的な行列分解の範囉を越えて、変換器ブロック内の埋め込み層およびすべての線形層に適用すること。
• 教師モデルの中間層出力を用いた2段階の知識蒸留戦略により、圧縮モデルの性能を向上させること。
• 標準的な自然言語処理ベンチマークにおいて、モデルサイズ、推論速度、精度のトレードオフを評価すること。

提案手法

• 埋め込み層、マルチヘッドアテンション、およびフィードフォワードネットワーク内の重み行列にKronecker分解を適用し、小さな行列のKronecker積として表現することで、パラメータ数を削減する。
• 各重み行列 W を W = A ⊗ B の形に分解し、A と B をより小さな行列として扱うことで、顕著なモデルサイズの縮小を実現する。
• 知識蒸留を事前学習段階および微調整段階の両方で適用し、元の BERT_BASE を教師モデルとして、学生モデルの中間層表現をガイドする。
• フレームワークを用いて BERT_BASE を KroneckerBERT に圧縮し、知識蒸留の効果を検証するためのアブレーションスタディを実施する。
• スマートフォンおよびハイパフォーマンスCPU上で推論速度を評価し、圧縮による実世界の遅延低減を測定する。
• TorchScript を用いて圧縮モデルをシリアル化し、エッジデバイスでのデプロイと遅延測定を実施する。

実験結果

リサーチクエスチョン

• RQ1 Kronecker分解を変換器ベースのモデルの埋め込み層およびすべての線形層に効果的に適用することで、極端な圧縮を可能にすることができるか？
• RQ2 中間層知識蒸留は、Kronecker圧縮モデルにおける性能劣化をどのように緩和するか？
• RQ3 高い圧縮比（例：19倍）において、KroneckerBERT は標準的な自然言語処理ベンチマークでどの程度の性能を示すか？
• RQ4 Kronecker分解によるFLOPsの削減は、エッジデバイス上で実際に推論速度の向上に繋がるか？
• RQ5 分布外データにおける精度および耐性の観点から、KroneckerBERT は最先端の圧縮モデルと比較してどの程度優れているか？

主な発見

• KroneckerBERT は 19倍の圧縮比（BERT_BASEの5％）を達成し、この圧縮レベルにおいて既存の手法を顕著に上回っている。
• GLUEベンチマークでは、19倍圧縮の KroneckerBERT が最先端の圧縮 BERT モデルを上回り、圧縮と性能のトレードオフにおいて優れた性能を示している。
• SQuAD では、KroneckerBERT が最先端の圧縮手法を上回る優れた性能を示しており、強力なロバストネスと一般化能力を示している。
• アブレーションスタディにより、事前学習段階および微調整段階の両方で知識蒸留が不可欠であり、2段階のKDが最高の性能をもたらすことが確認された。
• エッジデバイス（スマートフォン）では、19倍圧縮の KroneckerBERT が最大で2.4倍の高速化を達成しており、実用的な推論利点が裏付けられている。
• ハイパフォーマンスCPUでは、Kronecker行列乗算の逐次的性質のため、速度向上が観察されなかった。これは、ハードウェアに配慮した最適化の重要性を示している。