[論文レビュー] Less is More: Task-aware Layer-wise Distillation for Language Model Compression

Liang Chen, Simiao Zuo+4 more|arXiv (Cornell University)|Oct 4, 2022

Topic Modeling被引用数 14

ひとこと要約

本稿では、教師モデルの隠れ表現から関連する知識のみを抽出するタスク固有のフィルタを用いて、学生モデルに知識蒸留を行う、タスクに配慮した階層別蒸留（TED）を提案する。この手法により、冗長性が低減され、蒸留効率が向上する。TEDは、継続的事前学習および微調整の両設定において最先端の性能を達成し、GLUEおよびSQuADベンチマークで既存の蒸留手法を上回る。

ABSTRACT

Layer-wise distillation is a powerful tool to compress large models (i.e. teacher models) into small ones (i.e., student models). The student distills knowledge from the teacher by mimicking the hidden representations of the teacher at every intermediate layer. However, layer-wise distillation is difficult. Since the student has a smaller model capacity than the teacher, it is often under-fitted. Furthermore, the hidden representations of the teacher contain redundant information that the student does not necessarily need for the target task's learning. To address these challenges, we propose a novel Task-aware layEr-wise Distillation (TED). TED designs task-aware filters to align the hidden representations of the student and the teacher at each layer. The filters select the knowledge that is useful for the target task from the hidden representations. As such, TED reduces the knowledge gap between the two models and helps the student to fit better on the target task. We evaluate TED in two scenarios: continual pre-training and fine-tuning. TED demonstrates significant and consistent improvements over existing distillation methods in both scenarios. Code is available at https://github.com/cliang1453/task-aware-distillation.

研究の動機と目的

• 階層別蒸留における知識の冗長性という課題に対処し、教師の隠れ状態に含まれる不要な情報のため、学生モデルが有用な表現を学習しづらくなる問題を解消すること。
• タスク関連の特徴に焦点を当てることで、大規模な教師モデルと小規模な学生モデルとの間の知識ギャップを縮小すること。
• 中間層からの予測性の高い知識にのみアライメントさせることで、リソースが限られた環境下での学生モデルの汎化性能と性能を向上させること。
• まずタスク関連の特徴を特定し、その後に効率的に蒸留を行う二段階の学習フレームワークを開発すること。
• 継続的事前学習および微調整の両設定において、既存の蒸留手法に比して一貫した性能向上を示すこと。

提案手法

• TEDは、学生モデルおよび教師モデルの各層に対してタスクに配慮したフィルタを導入し、タスク固有のヘッド（例：分類のためのソフトマックス）を備えたニューラルネットワークとして実装する。
• 第1段階では、教師および学生モデルのパラメータを固定した状態でフィルタを訓練し、隠れ表現からタスク関連の特徴を抽出する能力を学習する。この段階では、ターゲットタスクの損失を予測する。
• 第2段階では、フィルタを固定（タスクヘッドなし）し、対応する教師および学生層のフィルタリング出力間の平均二乗誤差（MSE）を最小化するように、学生モデルと同時に微調整を行う。
• 蒸留損失は、対応する教師および学生層のフィルタリング出力間の平均二乗誤差として定義され、タスク固有の知識の整合性を促進する。
• 全体の学習目的は、タスク固有の損失、予測蒸留損失、およびタスクに配慮した階層別蒸留損失を、学習可能な重みで組み合わせたものである。
• フィルタは線形射影、MLP、またはTransformer層として実装可能であり、アブレーションスタディにより性能のトレードオフが示されている。

実験結果

リサーチクエスチョン

• RQ1 中間隠れ表現のタスクに配慮したフィルタリングは、小規模な学生モデルにおける知識蒸留の効率を向上させるか？
• RQ2 教師の表現から冗長な情報を除去することで、標準的な階層別蒸留と比較して学生モデルの性能が向上するか？
• RQ3 継続的事前学習や微調整などの異なる学習パラダイムにおいて、TEDはどのように性能を発揮するか？
• RQ4 フィルタの事前学習とその後の蒸留という二段階の学習フレームワークは、エンドツーエンドの蒸留と比較して、より良い整合性と汎化性能をもたらすか？
• RQ5 異なるフィルタアーキテクチャ（例：線形 vs. MLP vs. Transformer）が蒸留性能に与える影響は何か？

主な発見

• TEDは、GLUEベンチマークにおいて、標準的な知識蒸留および既存の階層別蒸留手法を著しく上回り、特にデータが少ない状況下で顕著な向上を示す。
• SQuAD v1.1およびv2.0の質問応答データセットにおいて、DeBERTaV3-xsmallの学生モデルをDeBERTaV3-base教師モデルから蒸留することで、最先端の結果を達成した。
• 継続的事前学習の設定において、LambadaおよびWikiText-103のゼロショットおよびフェイントショットの転移性能が向上し、強力な汎化能力を示した。
• アブレーションスタディにより、タスクに配慮したフィルタが冗長な知識の悪影響を軽減し、蒸留効率と学生モデルの適合度を向上させることを示した。
• 二段階の学習プロセスは、フィルタと学生モデルを同時に学習するジョイントトレーニングよりも、収束性と性能が優れている。
• より複雑なフィルタ（例：MLPやTransformer）を用いることで、単純な線形射影よりも性能が向上するが、計算コストが増加する。