[論文レビュー] Knowledge Distillation from Internal Representations

Gustavo Aguilar, Ling Yuan+4 more|arXiv (Cornell University)|Oct 8, 2019

Topic Modeling参考文献 20被引用数 23

ひとこと要約

本稿では、BERT類似モデルの内部表現からの知識蒸留を提案し、大規模な教師モデルのソフトラベルおよび隠れ層活性化を模倣するように学生ネットワークを訓練する。注意マップと隠れ状態におけるKLダイバージェンスを用いて複数の内部層にわたり蒸留を実施することで、一般化性能が著しく向上し、GLUEベンチマークデータセット全体で標準的な知識蒸留を上回る。

ABSTRACT

Knowledge distillation is typically conducted by training a small model (the student) to mimic a large and cumbersome model (the teacher). The idea is to compress the knowledge from the teacher by using its output probabilities as soft-labels to optimize the student. However, when the teacher is considerably large, there is no guarantee that the internal knowledge of the teacher will be transferred into the student; even if the student closely matches the soft-labels, its internal representations may be considerably different. This internal mismatch can undermine the generalization capabilities originally intended to be transferred from the teacher to the student. In this paper, we propose to distill the internal representations of a large model such as BERT into a simplified version of it. We formulate two ways to distill such representations and various algorithms to conduct the distillation. We experiment with datasets from the GLUE benchmark and consistently show that adding knowledge distillation from internal representations is a more powerful method than only using soft-label distillation.

研究の動機と目的

• 標準的な知識蒸留には出力確率の移管しか行わず、教師の内部言語的知識を保持できないという限界を是正すること。
• 出力確率に加えて、大規模な教師モデルから小規模な学生モデルへ内部表現（例：注意マップや隠れ状態）を移管することで、モデル圧縮を向上させること。
• 学生モデルの内部挙動を教師モデルと一致させることで、出力分布だけでなく、教師モデルと同様に一般化できるようにすること。
• 特にリソースが限られた環境や制約のある推論設定において、ソフトラベル蒸留のみで達成可能な範囲を超えて学生の性能を向上させる手法を開発すること。

提案手法

• 本手法は、最終出力確率に加え、教師モデルの複数層にわたる中間隠れ表現に対しても知識蒸留を実施する。
• 各層で教師モデルと学生モデルの注意確率を一致させるためにKLダイバージェンスを用い、高レベルの言語的抽象化を捉える。
• 特定の層において、教師モデルと学生モデルの隠れ状態の間で別個のKLダイバージェンス損失を適用することで、内部表現の構造的整合性を実現する。
• プログレッシブおよびスタックド知識蒸留戦略をサポートし、複数の教師層を1つの学生層に柔軟に圧縮可能である。
• 学生はハードラベルの交差エントロピー損失、ソフトラベルの蒸留損失に加え、注意マップと隠れ状態の追加蒸留損失を組み合わせて訓練される。
• FitNetのような先行手法とは異なり、表現整合性のための追加パラメータを導入しない。代わりに、同じ空間内での直接的な表現一致を実現する。

実験結果

リサーチクエスチョン

• RQ1 大規模な教師モデルの内部表現（例：注意マップや隠れ状態）を蒸留することで、標準的な知識蒸留を超えて、小規模な学生モデルの一般化性能と性能が向上するか？
• RQ2 出力確率の一致を超えて、学生の内部挙動を教師モデルと一致させることで、曖昧または困難な例においてもより頑健で正確な予測が得られるか？
• RQ3 誤差のパターンや教師の予測との一貫性の観点から、内部表現蒸留は標準KDと比べてどのように異なるか？
• RQ4 追加のパラメータを導入せず、推論の複雑さを増さずに、深層トランスフォーマー層に符号化された言語的知識を効果的に圧縮できるか？
• RQ5 早期層と後期層で蒸留を実施する場合の違い（例：早期 vs. 後期）が、最終的な学生モデルの性能および一般化能力に与える影響は何か？

主な発見

• 内部表現蒸留を施した学生モデルは、QQP開発セットにおいて教師の正解予測の97.9％（36,191／36,967）を達成した。一方、標準KDでは95.7％（35,401）であった。
• 内部表現蒸留を施した学生は776件の誤予測をしたが、標準KDの学生は1,566件であった。これは、教師の意思決定境界とより良好に一致していることを示している。
• QQPデータセットでは、教師が正解であったが標準KDの学生が失敗したケース（表4のサンプル3および4）において、内部表現蒸留を施した学生は100％の正答率を達成した。
• 教師が正解であった場合、内部表現蒸留により誤予測数が標準KDと比べて50％以上減少した。これは、より優れた一貫性と一般化性能を示している。
• 教師が誤りを犯した場合（例：表4のサンプル1）、内部表現蒸留を施した学生はより慎重な予測（確率0.4221）を示したのに対し、標準KDの学生は0.9999という過信度の高い予測を示した。これは、過信度の低い、より良好にキャリブレーションされた予測であることを示唆している。
• 本手法は、すべての4つのGLUEベンチマークデータセットで標準KDを一貫して上回り、精度および教師モデルの挙動との整合性の両面で顕著な向上を示した。