[論文レビュー] ALP-KD: Attention-Based Layer Projection for Knowledge Distillation

Peyman Passban, Yimeng Wu+2 more|arXiv (Cornell University)|Dec 27, 2020

Topic Modeling参考文献 29被引用数 8

ひとこと要約

本稿では、学習可能な注目重みを用いて複数の教師層を組み合わせる注目ベースのレイヤープロジェクション手法ALP-KDを提案する。この手法により、層のスキップを回避し、最終層および中間層の両方の表現を模倣できるため、BERTをより小さなモデルに distillation する際、GLUEベンチマークで既存手法を上回る優れた性能を達成する。

ABSTRACT

Knowledge distillation is considered as a training and compression strategy in which two neural networks, namely a teacher and a student, are coupled together during training. The teacher network is supposed to be a trustworthy predictor and the student tries to mimic its predictions. Usually, a student with a lighter architecture is selected so we can achieve compression and yet deliver high-quality results. In such a setting, distillation only happens for final predictions whereas the student could also benefit from teacher's supervision for internal components. Motivated by this, we studied the problem of distillation for intermediate layers. Since there might not be a one-to-one alignment between student and teacher layers, existing techniques skip some teacher layers and only distill from a subset of them. This shortcoming directly impacts quality, so we instead propose a combinatorial technique which relies on attention. Our model fuses teacher-side information and takes each layer's significance into consideration, then performs distillation between combined teacher layers and those of the student. Using our technique, we distilled a 12-layer BERT (Devlin et al. 2019) into 6-, 4-, and 2-layer counterparts and evaluated them on GLUE tasks (Wang et al. 2018). Experimental results show that our combinatorial approach is able to outperform other existing techniques.

研究の動機と目的

• アーキテクチャの不一致により中間教師層をスキップする既存の知識蒸留手法の制限を解消すること。
• 一対一のレイヤーマッチングに依存せず、非同一の複数の教師層からの監視を可能にすることで学生の性能を向上させること。
• 動的かつ組み合わせ的に関連する教師層表現を重み付け統合する、組み合わせ的で注目ベースのメカニズムを開発すること。
• 注目による教師層の統合が、学生モデルと教師モデル間の内部表現の整合性を向上させることを実証すること。
• 深層教師ネットワーク（例：12層Bert）をコンactな学生モデル（6層、4層、2層バージョン）に蒸留する際、最先端の性能を達成すること。

提案手法

• 蒸留用に複数の教師層特徴を1つの重み付き組み合わせに統合する、新規の注目ベースのレイヤープロジェクション（ALP-KD）メカニズムを提案する。
• 学生の現在のレイヤーに最も関連する教師層の重要度を、学習可能な注目メカニズムを用いて割り当てる。
• 複数の教師層からの中間特徴の重み付き和として、統合された教師表現を定式化する。重みはマルチヘッド注目モジュールによって計算される。
• 学生の中間出力と、教師特徴の注目重み付き組み合わせとの間の差分を比較する蒸留損失を用い、統合された教師表現を学生の学習プロセスに統合する。
• 隣接しない、またはアーキテクチャが一致しない層に対しても注目を許容することで、柔軟な構成が可能になり、レイヤー単位の整合性が不要になる。
• ハードラベルの交差エントロピー損失と、学生出力と注目重み付き教師特徴との間の蒸留損失を組み合わせた損失関数を用いて学生モデルを訓練する。

実験結果

リサーチクエスチョン

• RQ1 注目を用いて非同一の複数の教師層を統合することで、一対一のレイヤーマッチングを超える知識蒸留の向上が達成できるか？
• RQ2 教師表現の注目ベース統合は、学生と教師モデル間の中間隠れ状態の整合性にどのように影響するか？
• RQ3 アーキテクチャの不一致による層のスキップを回避することで、学生モデルの性能に顕著な向上が見られるか？
• RQ4 NLPベンチマークにおいて、PKD や RKD といった既存手法に比べて、ALP-KD は正確性および表現の忠実度の面でどの程度優れているか？
• RQ5 12層Bertのような深層モデルを、2層Bertのような大幅に小型化されたアーキテクチャに蒸留する際、ALP-KDは性能劣化を伴わずに効果的に適用可能か？

主な発見

• 12層Bertを6層、4層、2層の学生モデルに蒸留する際、ALP-KDはPKD や RKD といった既存の知識蒸留手法を、複数のGLUEベンチマークタスクで上回る性能を達成する。
• 注目ベースの教師層統合により、学生と教師モデル間の中間表現が著しく近くなる。コサイン距離の可視化では、ALP-KDは平均距離0.05を達成したのに対し、PKDは0.20であった。
• PCAを用いた隠れ状態の可視化により、ALP-KDの学生モデルが、ベースライン手法よりも教師の表現に近いものを学習していることが確認された。
• 教師層のすべての情報を活用することで、アーキテクチャの不一致による層のスキップ問題を効果的に緩和した。
• ALP-KDは、非常に深い教師モデルから非常に狭い学生モデルに蒸留する場合でも、アーキテクチャの乖離に対しても強固な性能を維持している。
• アブレーションスタディにより、注目ベースの組み合わせが不可欠であることが確認された。単純な連結や均等な重み付けでは、競争力のある結果は得られなかった。