A Dual-Space Framework for General Knowledge Distillation of Large Language Models

要約

知識蒸留（KD）は、知識を小さなモデルに転送することにより、大規模な言語モデル（LLM）を圧縮するための有望なソリューションです。
このプロセス中、ホワイトボックスKDメソッドは通常、教師モデルの出力分布と学生モデルの間の距離を最小限に抑えて、より多くの情報を転送します。
ただし、現在のホワイトボックスKDフレームワークが2つの制限を示していることを明らかにします。a）異なる出力スペースからの橋渡し確率分布は、教師モデルと学生モデルの類似性を制限します。
b）このフレームワークは、異なる語彙を持つLLMに適用することはできません。
これらの制限の根本原因の1つは、KDの教師と生徒からの分布が異なる予測ヘッドによって出力され、異なる出力スペースと寸法に分布が生成されることです。
したがって、この論文では、教師の予測ヘッドとKDの学生モデルを統一するデュアルスペースの知識蒸留（DSKD）フレームワークを提案します。
具体的には、最初に2つのプロジェクターを導入して、教師/生徒の隠された状態を学生/教師の表現スペースに投影するための理想的な初期化を行います。
この後、異なるモデルの隠された状態は同じヘッドを共有し、分布の出力スペースを統合することができます。
さらに、2つの異なるトークン化されたシーケンスで同じトークンを整列させるために、正確なトークンアライメント（ETA）アルゴリズムを開発します。
上記に基づいて、当社のDSKDフレームワークは、ポリティとポリティのKDの両方をサポートする一般的なKDフレームワークであり、語彙に関係なく任意の2つのLLM間のKDをサポートしています。
指導、数学的推論、およびコード生成ベンチマークに関する広範な実験は、DSKDが現在のホワイトボックスKDフレームワークに基づいて既存の方法を大幅に上回り、LLMSの他のクロストコーネザーKDメソッドを異なる語彙で上回ることを示しています。

要約(オリジナル)

Knowledge distillation (KD) is a promising solution to compress large language models (LLMs) by transferring their knowledge to smaller models. During this process, white-box KD methods usually minimize the distance between the output distributions of the teacher model and the student model to transfer more information. However, we reveal that the current white-box KD framework exhibits two limitations: a) bridging probability distributions from different output spaces will limit the similarity between the teacher model and the student model; b) this framework cannot be applied to LLMs with different vocabularies. One of the root causes for these limitations is that the distributions from the teacher and the student for KD are output by different prediction heads, which yield distributions in different output spaces and dimensions. Therefore, in this paper, we propose a dual-space knowledge distillation (DSKD) framework that unifies the prediction heads of the teacher and the student models for KD. Specifically, we first introduce two projectors with ideal initialization to project the teacher/student hidden states into the student/teacher representation spaces. After this, the hidden states from different models can share the same head and unify the output spaces of the distributions. Furthermore, we develop an exact token alignment (ETA) algorithm to align the same tokens in two differently-tokenized sequences. Based on the above, our DSKD framework is a general KD framework that supports both off-policy and on-policy KD, and KD between any two LLMs regardless of their vocabularies. Extensive experiments on instruction-following, mathematical reasoning, and code generation benchmarks show that DSKD significantly outperforms existing methods based on the current white-box KD framework and surpasses other cross-tokenizer KD methods for LLMs with different vocabularies.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google