[論文レビュー] LRC-BERT: Latent-representation Contrastive Knowledge Distillation for Natural Language Understanding

Hao Fu, Shaojun Zhou+5 more|arXiv (Cornell University)|Dec 14, 2020

Topic Modeling参考文献 40被引用数 5

ひとこと要約

LRC-BERTは、中間層の表現を転送するための新しい知識蒸留フレームワークを提案する。このフレームワークは、角距離に基づく対照学習（COS-NCE）を用い、出力値マッチングを超えた構造的分布特徴を捉えることで、知識蒸留を強化する。さらに勾配摂動を導入することでモデルの頑健性を向上させ、2段階訓練戦略を採用することで、GLUEベンチマークで最先端の性能を達成し、一般化性能と安定性が向上した。

ABSTRACT

The pre-training models such as BERT have achieved great results in various natural language processing problems. However, a large number of parameters need significant amounts of memory and the consumption of inference time, which makes it difficult to deploy them on edge devices. In this work, we propose a knowledge distillation method LRC-BERT based on contrastive learning to fit the output of the intermediate layer from the angular distance aspect, which is not considered by the existing distillation methods. Furthermore, we introduce a gradient perturbation-based training architecture in the training phase to increase the robustness of LRC-BERT, which is the first attempt in knowledge distillation. Additionally, in order to better capture the distribution characteristics of the intermediate layer, we design a two-stage training method for the total distillation loss. Finally, by verifying 8 datasets on the General Language Understanding Evaluation (GLUE) benchmark, the performance of the proposed LRC-BERT exceeds the existing state-of-the-art methods, which proves the effectiveness of our method.

研究の動機と目的

• 従来のNLP分野における知識蒸留手法の限界に取り組む。これらの手法は出力値マッチングにのみ焦点を当てており、中間層表現の構造的分布特性を無視している。
• 知識蒸留の訓練中にモデルの頑健性を向上させるために、勾配摂動に基づくアーキテクチャを導入する。これは知識蒸留分野で初めての試みである。
• 教師モデルから学生モデルへの知識転送を強化するため、初期段階で分布構造学習を優先する2段階訓練戦略を設計する。
• 学生モデルのアーキテクチャを教師の層構造から分離することで、より柔軟な学生モデル設計を可能にし、圧縮の柔軟性を高める。

提案手法

• 中間層表現間の角距離類似度を捉えるために、コサインベースのノイズ対照推定に基づく新しい対照損失、COS-NCEを導入する。これにより、構造的分布特徴をモデル化できる。
• 全蒸留損失を段階的に最適化する2段階訓練戦略を採用する。初期段階では中間層の分布学習に焦点を当て、その後で予測ヘッド損失を統合する。
• 訓練中に単語埋め込み層に対して勾配摂動を適用する。勾配値を用いて隠れ状態を摂動させることで、モデルの頑健性を向上させる。
• 複数の損失成分を組み合わせる：ソフトラベル用のKLダイバージェンス、ハードラベル用の交差エントロピー、中間層蒸留用のCOS-NCE。これらは適応的重み付けが可能である。
• 次元変換行列を用いて、学生ネットワークの設計を柔軟に可能にし、学生アーキテクチャを教師の層ごとの正確な対応から分離する。
• 教師よりパラメータ数が少ない学生モデルを採用することで、エッジデバイスへの効率的なデプロイが可能でありながら、高い性能を維持できる。

実験結果

リサーチクエスチョン

• RQ1 角距離に基づく対照学習は、中間表現の値マッチングを超えて、NLP分野における知識蒸留を改善できるか？
• RQ2 訓練中に勾配摂動を適用することで、知識蒸留における学生モデルの頑健性が向上するか？
• RQ3 2段階訓練戦略により、初期段階で分布構造学習を優先することで、蒸留の有効性が向上するか？
• RQ4 標準ベンチマークにおいて、提案手法は従来の蒸留手法と比較して、精度、頑健性、圧縮効率の観点でどのように異なるか？

主な発見

• LRC-BERTはGLUEベンチマークの8つのデータセットすべてで最先端の性能を達成し、従来の蒸留手法を上回った。
• アブレーションスタディの結果、COS-NCE損失を削除すると、MNLI-m開発セットで4.0%の精度低下が生じ、これが知識転送において極めて重要な役割を果たしていることが確認された。
• 2段階訓練法は、非ステージドバージョンと比較して4.0%の性能向上を示し、初期段階での分布構造学習の優先が蒸留効果を高めることを示した。
• 勾配摂動により、2段階目の訓練損失が安定化し、フラクチュエーションが減少し、モデルの収束性と頑健性が向上した。
• COS-NCE損失は角距離類似度特徴を効果的に捉えており、MSEベースのベースラインが大きな埋め込み差異のため失敗する状況でも正しく予測を可能にした。
• CoLAタスクでは、ベースライン蒸留手法と比較して13%の改善を示し、角距離に基づく蒸留の利点が顕著に現れた。