[論文レビュー] Pre-trained Summarization Distillation

Sam Shleifer, Alexander M. Rush|arXiv (Cornell University)|Oct 24, 2020

Topic Modeling参考文献 25被引用数 57

ひとこと要約

本論文は、shrinking-and-fine-tune (SFT)、知識蒸留 (KD)、疑似ラベリング (PL) を比較し、巨大な事前学習要約モデル（BART、Pegasus）をより小さく高速な student に蒸留する。CNN/DailyMail ではSFTがしばしば優れる一方、XSUM ではPLが勝つか同等となるなど、トレードオフはさまざま。

ABSTRACT

Recent state-of-the-art approaches to summarization utilize large pre-trained Transformer models. Distilling these models to smaller student models has become critically important for practical use; however there are many different distillation methods proposed by the NLP literature. Recent work on distilling BERT for classification and regression tasks shows strong performance using direct knowledge distillation. Alternatively, machine translation practitioners distill using pseudo-labeling, where a small model is trained on the translations of a larger model. A third, simpler approach is to 'shrink and fine-tune' (SFT), which avoids any explicit distillation by copying parameters to a smaller student model and then fine-tuning. We compare these three approaches for distillation of Pegasus and BART, the current and former state of the art, pre-trained summarization models, and find that SFT outperforms knowledge distillation and pseudo-labeling on the CNN/DailyMail dataset, but under-performs pseudo-labeling on the more abstractive XSUM dataset. PyTorch Code and checkpoints of different sizes are available through Hugging Face transformers here http://tiny.cc/4iy0tz.

研究の動機と目的

• 実用のために大規模な事前学習済み要約モデルの圧縮を動機づける。
• 標準データセット CNN/Daily Mail および XSUM で3つの蒸留系（SFT、KD、PL）を体系的に比較する。
• タスクとモデルサイズに基づく蒸留戦略の選択について、実装の詳細と実証的な指針を提供する。
• 品質（ROUGE）と速度（推論時間と訓練時間）のトレードオフを定量化する。

提案手法

• 3つの蒸留パラダイムを定義する：Shrink and Fine-Tune（SFT）、Pseudo-Labeling（PL）、直接的なKnowledge Distillation（KD）。
• 教師モデルから最大に間引かれた層をコピーして学生モデルを初期化し、タスクデータでファインチューニングする。
• KD では、ロジット損失、データ損失、隠れ状態損失の加重和を最小化する（L_KD = alpha_logits L_logits + alpha_data L_data + alpha_hidden L_hidden）。
• ビームサーチで教師の要約を生成し、これらの疑似ターゲットで学生を訓練する（L_pseudo）。
• CNN/Daily Mail（要約はより控えめ）とXSUM（より抽象的）で評価し、複数の教師-学生ペアに対して ROUGE-1/2/L を報告する。
• 各蒸留手法の訓練時間コストと推論速度の向上を報告する。

実験結果

リサーチクエスチョン

• RQ1 BART および Pegasus に対して、どの蒸留手法が速度と要約品質の最良のトレードオフをもたらすか？
• RQ2 教師から最大に間引かれた層をコピーして初期化することは、事前学習なしで有効な圧縮を可能にするか？
• RQ3 CNN/Daily Mail と XSUM（抽象的要約）では蒸留手法の性能はどう異なるか？
• RQ4 疑似ラベルの品質とタスク横断転移（翻訳ベンチマーク）ある蒸留性能へ与える影響は？

主な発見

教師	サイズ	データ	教師スコア	SFTスコア	KDスコア	疑似スコア	コスト
BART	12-3	XSUM	22.29	21.08	21.63	21.38	6
Pegasus	16-4	XSUM	24.56	22.64	21.92	23.18	34
BART	12-6	CNN	21.06	21.21	20.95	19.93	19.5
Pegasus	16-4	CNN	21.37	21.29	-	20.1	48
Marian	6-3	EN-RO	27.69	25.91	24.96	26.85	28
mBART	12-3	EN-RO	26.46	25.61	25.87	26.09	50

• CNN では、SFT が BART・Pegasus を含む教員全体で、より高価な KD および PL を上回る。
• XSUM では、BART についてはKD、Pegasus についてはPL が SFT を上回る場合があり、PL は特定の設定で教員の性能に最も近い。
• データセットとモデルペアに依存するが、一部の蒸留モデルは教師と比べてROUGEの損失は小さく、速度を著しく向上（例: 43-93% 高速化）させている。
• 疑似ラベルの品質が高く、元のファインチューニングデータと組み合わせた場合、疑似ラベルは性能を大幅に向上させる（Orig+PL、Orig+PL+PL）。
• XSUM に関しては KD は SFT に比べて限定的な利得で、教師の順伝播が必要なため計算コストが高い。PL は特に Pegasus にとって強力な代替手段になり得る。