[論文レビュー] DistiLLM: Towards Streamlined Distillation for Large Language Models

Jongwoo Ko, Sungnyun Kim+2 more|arXiv (Cornell University)|Feb 6, 2024

Topic Modeling被引用数 4

ひとこと要約

DistiLLM は、自己回帰的大型言語モデル向けに洗練された知識蒸留フレームワークを提案し、安定的で収束性の高いトレーニングを実現するための新しい歪みKullback-Leibler発散損失関数と、効率的に学生モデルが出力する結果を活用するための適応的オフポリシー手法を導入している。この手法は、指示従い型および要約タスクにおいて最先端の性能を達成し、最近のKD手法と比較して最大4.3倍の高速トレーニングを実現している。

ABSTRACT

Knowledge distillation (KD) is widely used for compressing a teacher model to a smaller student model, reducing its inference cost and memory footprint while preserving model capabilities. However, current KD methods for auto-regressive sequence models (e.g., large language models) suffer from missing a standardized objective function. Moreover, the recent use of student-generated outputs to address training-inference mismatches has significantly escalated computational costs. To tackle these issues, we introduce DistiLLM, a more effective and efficient KD framework for auto-regressive language models. DistiLLM comprises two components: (1) a novel skew Kullback-Leibler divergence loss, where we unveil and leverage its theoretical properties, and (2) an adaptive off-policy approach designed to enhance the efficiency in utilizing student-generated outputs. Extensive experiments, including instruction-following tasks, demonstrate the effectiveness of DistiLLM in building high-performing student models while achieving up to 4.3$ imes$ speedup compared to recent KD methods.

研究の動機と目的

• 自己回帰的言語モデルにおける知識蒸留の分野で、標準的で安定した目的関数の欠如を是正すること。
• オフポリシー蒸留における学生生成出力（SGO）の過剰生成が引き起こすトレーニング非効率性を低減すること。
• 学生モデルの出力におけるモード崩壊と分布不一致を最小限に抑えることで、蒸留の効果を高めること。
• 大幅に短縮されたトレーニング時間と計算コストで高性能な学生モデルを達成すること。
• 理論的裏付けが強く、効率的かつ効果的な大規模言語モデルの蒸留フレームワークを提供すること。

提案手法

• 勾配の安定化と近似誤差の低減を図る強固な理論的基盤を持つ歪みKullback-Leibler（KLD）発散損失関数を導入し、蒸留の過程で使用する。
• 学生生成出力の使用割合を動的に制御する適応的オフポリシー戦略を採用し、サンプル効率を向上させる。
• 固定データセットの監視とオンポリシーの学生生成シーケンスの両方を活用し、分布の整合性とトレーニング効率のバランスを取る。
• 教師と学生の出力分布間の非対称発散を最小化するように蒸留目的関数を最適化し、モード崩壊を回避する。
• 段階的に学生生成出力への依存度を高めるカリキュラムのようなトレーニングスケジュールを適用し、学習の安定化を図る。
• シーケンスレベルの出力に歪みKLD損失を適用することで、多様な自己回帰的タスクにおいて教師の生成行動とより良い一致を実現する。

Figure 1 : Examples of SGOs from GPT-2 (student) and their corresponding validation loss by GPT-2 XL (teacher). Since the teacher model may be not familiar with the SGO, using $p(\mathbf{y}|\mathbf{x})$ as a target distribution can misguide the student model, as shown in Tab. 7 .

実験結果

リサーチクエスチョン

• RQ1 理論的裏付けのある歪み対称KLD損失は、自己回帰的LMの知識蒸留におけるトレーニングの安定性と収束性を向上させることができるか？
• RQ2 高い計算コストを伴わずに、学生生成出力を効率的かつ適応的に蒸留に組み込む方法は何か？
• RQ3 固定データセットの監視と適応的オフポリシーSGOを組み合わせることで、オンポリシーまたは固定SGO戦略と比較して、より優れた一般化性能と収束速度が得られるか？
• RQ4 提案手法は、生成タスクにおける学生モデルの性能を維持または向上させつつ、トレーニング時間をどの程度短縮できるか？
• RQ5 歪みKLD損失は、標準的KLDおよび逆方向KLDアプローチと比較して、多様な指示従い型およびテキスト生成ベンチマークで頑健であるか？

主な発見

• DistiLLM は、指示従い型およびテキスト要約タスクにおいて、既存のKD手法を上回る最先端の性能を達成した。
• 最近のオフポリシーKD手法（Gu et al. (2024) や Agarwal et al. (2024)）と比較して、トレーニング時間を最大4.3倍短縮した。
• 歪みKLD損失により、収束が速く、トレーニングがより安定した。標準KLDで観察された勾配爆発のリスクも低減された。
• 適応的オフポリシーモジュールにより、学生生成出力の使用を動的に調整することで、無駄や低品質なシーケンスの使用を回避し、サンプル効率が向上した。
• Self-Instruct および Vicuna ベンチマークにおいて、他の手法が失敗したり幻覚を含む出力を生成するケースでも、DistiLLM は100％の確率で簡潔で正確な応答を正しく生成した。
• 限られたSGO使用量でも優れた性能を発揮したため、高品質な蒸留を継続的かつ高価な学生生成を必要とせずに実現可能であることが示された。

Figure 2 : ( Left ): Normalized runtime according to the maximum response length of SGOs with GPT-2 XL teacher and GPT-2 student. ( Right ): Normalized runtime for various sizes of teacher and student models with a response length of 256. FWD and BWD denote forward and backward propagation, respecti