[論文レビュー] Contrastive Distillation on Intermediate Representations for Language Model Compression

Siqi Sun, Zhe Gan+4 more|arXiv (Cornell University)|Sep 29, 2020

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

ひとこと要約

本稿では、中間表現を対照的目的関数を用いて活用することで、圧縮言語モデルにおける知識蒸留を向上させるコントラスト的蒸留フレームワークCoDIRを提案する。GLUEベンチマークにおいて、より小型で高速なモデルで、既存手法を上回る性能を達成した。

ABSTRACT

Existing language model compression methods mostly use a simple L2 loss to distill knowledge in the intermediate representations of a large BERT model to a smaller one. Although widely used, this objective by design assumes that all the dimensions of hidden representations are independent, failing to capture important structural knowledge in the intermediate layers of the teacher network. To achieve better distillation efficacy, we propose Contrastive Distillation on Intermediate Representations (CoDIR), a principled knowledge distillation framework where the student is trained to distill knowledge through intermediate layers of the teacher via a contrastive objective. By learning to distinguish positive sample from a large set of negative samples, CoDIR facilitates the student's exploitation of rich information in teacher's hidden layers. CoDIR can be readily applied to compress large-scale language models in both pre-training and finetuning stages, and achieves superb performance on the GLUE benchmark, outperforming state-of-the-art compression methods.

研究の動機と目的

• BERTの中間層における構造的依存関係を捉えられない、独立した$L_2$損失に依存する従来の知識蒸留の限界を解消すること。
• 大規模言語モデルにおける知識蒸留にコントラスト学習を適用し、中間表現における高階の関係を活用すること。
• ラベルが入手不能または希少なPre-trainingおよびFine-tuningの両ステージで効果的なサンプリング戦略を設計すること。
• 標準的な$L_2$蒸留を対照的目的関数に置き換えることで、同一サンプル（一致）と異なるサンプル（不一致）の表現ペアを区別するように学生モデルを訓練することにより、モデルの効率性と性能を向上させること。
• 平均プーリング表現が[CLS]トークンよりも、特に微調整段階で蒸留に有効であることを示すこと。

提案手法

• CoDIRは、教師および学生ネットワークの中間隠れ表現にコントラスト損失を適用し、学生モデルが同一入力からのペア（一致ペア）と多数の異なる入力からのペア（不一致ペア）を区別できるように訓練する。
• コントラスト的目的関数は、同一サンプルからの表現（ポジティブペア）を引き寄せ、異なるサンプルからの表現（ネガティブペア）を遠ざけるように定式化される。
• 現在のバッチからのすべての表現を格納するメモリバンクを用意し、特にラベルが入手できないPre-training段階においても、効率的なネガティブサンプリングを可能にする。
• Fine-tuningでは、異なるラベルを持つサンプルを用いて不一致ペアを構築し、類似度を最大限に低下させる。Pre-trainingでは、同じミニバッチからのランダムペアを代理として用いる。
• Pre-training段階では、コントラスト的蒸留と標準的な知識蒸留（ログイットにおけるKLダイバージェンス）およびマスク言語モデル（MLM）損失を併用する。
• 蒸留のための文レベル埋め込みとして、すべてのトークン表現の平均プーリングを用い、[CLS]トークンを置き換える。実験的に、これがより効果的であることが示された。

実験結果

リサーチクエスチョン

• RQ1 標準的な$L_2$ベースの蒸留と比較して、中間表現におけるコントラスト学習が、圧縮言語モデルにおける知識蒸留を向上させられるか？
• RQ2 Pre-training段階でラベルが入手できない状況下でも、Pre-trainingおよびFine-tuningの両段階で効果的なネガティブサンプリングを設計できるか？
• RQ3 [CLS]トークンの代わりに平均プールド表現を使用することで、圧縮モデルにおける蒸留性能が向上するか？
• RQ4 ネガティブサンプルの数が、言語モデル圧縮の文脈におけるコントラスト的蒸留の性能に与える影響は何か？
• RQ5 提案されたCoDIRフレームワークは、既存の圧縮手法と比較して、モデルサイズと推論時間を削減しながらGLUEで最先端の性能を達成できるか？

主な発見

• CoDIRはGLUEベンチマークにおいて、最先端の圧縮手法を上回り、標準的なMLM Pre-trainingに比べ1.9%の平均向上、KDオンリープリトレーニングに比べ1.0%の向上を達成した。
• BERT-baseの半分のサイズである学生モデルは、2倍の高速な推論速度と、GPU使用量の削減を実現しながら、競争力のある性能を示した。
• 平均プールド表現を蒸留ターゲットとして使用した場合、GLUE平均スコアは83.0に達し、[CLS]トークンベースの蒸留（平均スコア82.5）を上回った。
• ネガティブサンプル数を100から1000に増加させることで、大多数のGLUEタスクで性能が向上し、特にMNLIおよびQNLIで最大の向上が観察された。
• コントラスト損失は性能向上に顕著な寄与を示した：KDにCRD損失を追加すると、標準KDに比べ1.0%、バニラMLM Pre-trainingに比べ1.9%の平均GLUEスコア向上が得られた。
• モデル性能は、小さなデータセット（例：CoLA、MRPC、RTE）において高い分散を示しており、標準偏差が約1.5であった。これは、ランダム初期化のシードに敏感であることを示唆している。