[論文レビュー] BERT-EMD: Many-to-Many Layer Mapping for BERT Compression with Earth Mover's Distance

Jianquan Li, Xiaokang Liu+4 more|arXiv (Cornell University)|Oct 13, 2020

Topic Modeling参考文献 24被引用数 7

ひとこと要約

本論文では、教師モデルと生徒モデル間の多対多レイヤーマッピングを可能にする新規な BERT 蒸留手法 BERT-EMD を提案する。この手法は、知識伝達を最適化するために、地球移動距離（EMD）を用いる。コストアテンションメカニズムにより、適応的なレイヤー重みを学習することで、BERT-EMD は最先端の圧縮性能を達成し、パラメータ数と推論時間の両方を 50% に削減したにもかかわらず、GLUE タスクの 9 つのうち 7 つで BERT-Base を上回る性能を発揮する。

ABSTRACT

Pre-trained language models (e.g., BERT) have achieved significant success in various natural language processing (NLP) tasks. However, high storage and computational costs obstruct pre-trained language models to be effectively deployed on resource-constrained devices. In this paper, we propose a novel BERT distillation method based on many-to-many layer mapping, which allows each intermediate student layer to learn from any intermediate teacher layers. In this way, our model can learn from different teacher layers adaptively for various NLP tasks. %motivated by the intuition that different NLP tasks require different levels of linguistic knowledge contained in the intermediate layers of BERT. In addition, we leverage Earth Mover's Distance (EMD) to compute the minimum cumulative cost that must be paid to transform knowledge from teacher network to student network. EMD enables the effective matching for many-to-many layer mapping. %EMD can be applied to network layers with different sizes and effectively measures semantic distance between the teacher network and student network. Furthermore, we propose a cost attention mechanism to learn the layer weights used in EMD automatically, which is supposed to further improve the model's performance and accelerate convergence time. Extensive experiments on GLUE benchmark demonstrate that our model achieves competitive performance compared to strong competitors in terms of both accuracy and model compression.

研究の動機と目的

• 一対一のレイヤーマッピングに起因する BERT 蒸留の限界を是正すること。これは理論的根拠に欠け、多様な NLP タスクに適応できない。
• 各生徒レイヤーが、タスク固有の言語的知識要件に応じて、任意の中間教師レイヤーから動的に学習できるようにすること。
• 教師-生徒レイヤー間の距離を最適輸送問題として定式化し、地球移動距離（EMD）を用いて知識伝達を最適化すること。
• コストアテンションメカニズムにより、教師レイヤーの重要性を自動で学習し、性能向上と収束速度の加速を実現すること。
• 性能の損なわれない高い圧縮効率を実現し、リソース制約のあるデバイスへのデプロイを可能にすること。

提案手法

• 各中間生徒レイヤーがすべての中間教師レイヤーに注目できる多対多レイヤーマッピング戦略を提案。これにより、タスクに応じた適応的知識伝達が可能になる。
• 教師から生徒への知識の変換に必要な最小累積コストを計算するために、地球移動距離（EMD）を用い、隠れ表現の最適輸送をモデル化する。
• タスク固有のレイヤー重みを学習するコストアテンションメカニズムを導入。これにより、異なる教師レイヤーへの重要度が動的に割り当てられる。
• EMDに基づく知識蒸留を解くために輸送問題の定式化を採用。これにより、包括的かつ効率的な知識伝達が保証される。
• エンド・ツー・エンド最適化を実現するため、EMD に基づく知識蒸留損失とアテンション重み付きレイヤーマッチングを用いて生徒モデルを訓練する。
• 再現性のため、コードとデータを https://github.com/lxk00/BERT-EMD で公開する。

実験結果

リサーチクエスチョン

• RQ1 一対一マッピングと比較して、多対多レイヤーマッピング戦略は BERT 圧縮における知識蒸留を改善できるか？
• RQ2 地球移動距離（EMD）を距離指標として用いることで、BERT 蒸留における知識伝達とモデル性能が向上するか？
• RQ3 学習可能なコストアテンションメカニズムは、蒸留プロセスの適応性と収束速度を向上させられるか？
• RQ4 BERT-EMD は、多様な NLP タスクにおいて顕著に小型化されたモデルサイズと推論時間で、競争力のある性能を達成できるか？
• RQ5 言語的複雑性の異なるタスク、例えば RTE と MNLI において、モデルの性能はどのように変化するか？

主な発見

• 6 レイヤーの BERT-EMD は、GLUE タスクの 9 つ中 7 つで元の 12 レイヤー BERT-Base を上回り、強力な圧縮効率を示している。
• RTE タスクでは、BERT-Base よりも 5.3% の精度向上を達成しており、複雑な推論タスクにおける有効性が裏付けられている。
• STS-B タスクでは、Spearman相関係数で BERT-Base よりも 1% 向上しており、回帰的タスクにおける優れた性能を示している。
• アブレーションスタディの結果、EMD モジュールを削除すると、特に STS-B で性能が急激に低下することが判明。EMD が最適輸送の知識伝達において中心的な役割を果たしていることが裏付けられた。
• コストアテンションメカニズムは、性能向上と収束速度の加速を実現しており、タスク固有のレイヤー重要度を学習する価値があることが確認された。
• フロー行列の可視化により、異なるタスクが異なる教師レイヤーに注目していることが明らかになった。例えば、RTE では非対角マッピングが観察され、MNLI とは異なり、適応的かつタスクに特化した知識伝達が実現されている。