[論文レビュー] MiniLLM: On-Policy Distillation of Large Language Models

Yuxian Gu, Li Dong+2 more|arXiv (Cornell University)|Jun 14, 2023

Topic Modeling被引用数 19

ひとこと要約

MiniLLM は reverse Kullback–Leibler 発散を最小化して、より大きなホワイトボックス教師からより小さな生成系 LLM を蒸留し、モデルサイズを問わず命令の従い品質、校正、および長文生成を改善します。

ABSTRACT

Knowledge Distillation (KD) is a promising technique for reducing the high computational demand of large language models (LLMs). However, previous KD methods are primarily applied to white-box classification models or training small models to imitate black-box model APIs like ChatGPT. How to effectively distill the knowledge of white-box LLMs into small models is still under-explored, which becomes more important with the prosperity of open-source LLMs. In this work, we propose a KD approach that distills LLMs into smaller language models. We first replace the forward Kullback-Leibler divergence (KLD) objective in the standard KD approaches with reverse KLD, which is more suitable for KD on generative language models, to prevent the student model from overestimating the low-probability regions of the teacher distribution. Then, we derive an effective on-policy optimization approach to learn this objective. The student models are named MiniLLM. Extensive experiments in the instruction-following setting show that MiniLLM generates more precise responses with higher overall quality, lower exposure bias, better calibration, and higher long-text generation performance than the baselines. Our method is scalable for different model families with 120M to 13B parameters. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/minillm.

研究の動機と目的

• ホワイトボックス教師を用いたオープンで生成的な LLM の知識蒸留の動機付け。
• 低確率領域を過大評価しないよう reverse KLD に基づく最適化目的を提案。
• LLM の知識蒸留におけるポリシーグラデーション最適化を安定化させる実践的な訓練手法を開発。
• 多様なモデルファミリとサイズ（120M–13B）にわたる MiniLLM のスケーラビリティを実証。
• 命令従い品質、曝露バイアス、校正、長文生成の改善を評価。

提案手法

• 生成的 LLM の KD において前方 KLD を reverse KLD に置換。
• Policy Gradient を用いて min_theta KL[q_theta || p] の勾配を導出。
• 分散を低減する単一ステップ正則化を導入。
• 報酬改ざん（reward hacking）を緩和する教師混合サンプリングを実装。
• 生成長のバイアスを補正するための長さ正規化を適用。
• Pretraining データ上の言語モデル損失を含む PPO クリッピングで訓練。

実験結果

リサーチクエスチョン

• RQ1 reverse KLD ベースの蒸留（MiniLLM）が、ホワイトボックス教師が利用可能な場合に標準の forward-KLD KD より生成的 LLM に対して優れているか。
• RQ2 モデルファミリとサイズを跨いで、MiniLLM が曝露バイアス、校正、および長文生成品質にどう影響するか。
• RQ3 LLM のポリシーグラデーションベースの KD を安定化させるために必要な最適化戦略は何か。
• RQ4 MiniLLM は GPT-2、OPT、LLaMA ファミリ間で 120M から 13B パラメータへとスケールするか。
• RQ5 教師混合の強さと事前訓練の損失が最終的な性能とロバスト性に与える影響は何か。

主な発見

• MiniLLM は複数のデータセットとモデルサイズで標準的な KD のベースライン（語彙レベル KD およびシーケンスレベル KD）を一貫して上回る。
• Rouge-L スコアは、多くの設定で生成の精度と真の ground-truth との整合性が高いことを示す。
• MiniLLM は KD および SeqKD ベースラインに比べて曝露バイアスが低く、校正が改善される。
• MiniLLM を用いた蒸留モデルは長文生成性能が強く、生成多様性を維持する（4-gram の分離と LM 損失は安定したまま）。
• MiniLLM からの改善はモデルファミリ（GPT-2、OPT、LLaMA）およびパラメータスケール（120M–13B）を跨って安定している。
• アブレーションは、教師混合サンプリング、長さ正規化、単一ステップ正則化が安定性と性能のために重要であることを示している。