Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Dual-Space Knowledge Distillation for Large Language Models

Songming Zhang, Xue Zhang, Zengkui Sun, Yufeng Chen, Jinan Xu

TL;DR

Experiments on task-agnostic instruction-following benchmarks show that DSKD significantly outperforms the current white-box KD framework with various distance functions, and also surpasses existing KD methods for LLMs with different vocabularies.

Abstract

Knowledge distillation (KD) is known as a promising solution to compress large language models (LLMs) via transferring their knowledge to smaller models. During this process, white-box KD methods usually minimize the distance between the output distributions of the two models so that more knowledge can be transferred. However, in the current white-box KD framework, the output distributions are from the respective output spaces of the two models, using their own prediction heads. We argue that the space discrepancy will lead to low similarity between the teacher model and the student model on both representation and distribution levels. Furthermore, this discrepancy also hinders the KD process between models with different vocabularies, which is common for current LLMs. To address these issues, we propose a dual-space knowledge distillation (DSKD) framework that unifies the output spaces of the two models for KD. On the basis of DSKD, we further develop a cross-model attention mechanism, which can automatically align the representations of the two models with different vocabularies. Thus, our framework is not only compatible with various distance functions for KD (e.g., KL divergence) like the current framework, but also supports KD between any two LLMs regardless of their vocabularies. Experiments on task-agnostic instruction-following benchmarks show that DSKD significantly outperforms the current white-box KD framework with various distance functions, and also surpasses existing KD methods for LLMs with different vocabularies.

Dual-Space Knowledge Distillation for Large Language Models

TL;DR

Experiments on task-agnostic instruction-following benchmarks show that DSKD significantly outperforms the current white-box KD framework with various distance functions, and also surpasses existing KD methods for LLMs with different vocabularies.

Abstract

Knowledge distillation (KD) is known as a promising solution to compress large language models (LLMs) via transferring their knowledge to smaller models. During this process, white-box KD methods usually minimize the distance between the output distributions of the two models so that more knowledge can be transferred. However, in the current white-box KD framework, the output distributions are from the respective output spaces of the two models, using their own prediction heads. We argue that the space discrepancy will lead to low similarity between the teacher model and the student model on both representation and distribution levels. Furthermore, this discrepancy also hinders the KD process between models with different vocabularies, which is common for current LLMs. To address these issues, we propose a dual-space knowledge distillation (DSKD) framework that unifies the output spaces of the two models for KD. On the basis of DSKD, we further develop a cross-model attention mechanism, which can automatically align the representations of the two models with different vocabularies. Thus, our framework is not only compatible with various distance functions for KD (e.g., KL divergence) like the current framework, but also supports KD between any two LLMs regardless of their vocabularies. Experiments on task-agnostic instruction-following benchmarks show that DSKD significantly outperforms the current white-box KD framework with various distance functions, and also surpasses existing KD methods for LLMs with different vocabularies.
Paper Structure (41 sections, 21 equations, 11 figures, 6 tables)

This paper contains 41 sections, 21 equations, 11 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Background and Preliminary Study
  3. Current Framework for White-Box KD
  4. Limitations of the Current Framework
  5. Low Teacher-Student Similarity
  6. Dependency on the Same Vocabulary
  7. Methodology
  8. Dual-Space Knowledge Distillation Framework
  9. KD in Student Space.
  10. KD in Teacher Space.
  11. Cross-Model Attention Mechanism
  12. Experiments
  13. Experimental Setup
  14. Data.
  15. Models.
  16. ...and 26 more sections

Figures (11)

  • Figure 1: Simulation results with KL divergence as the distance function $\mathcal{D}(\cdot||\cdot)$. (a), (b) and (c) plot the student's hidden states and the teacher's hidden states before and after the two KD processes. (d) shows the convergence curves of $\mathcal{L}_{kd}$ in the two KD processes.
  • Figure 2: Win rates (%) on the response quality between TinyLLaMA trained by DSKD and the current white-box KD framework.
  • Figure 3: Distance between the representation structures of the teacher and the student.
  • Figure 4: Simulation results with reverse KL divergence as the KD objective. (a), (b) and (c) plot the student's hidden states and the teacher's hidden states before and after the two KD processes. (d) shows the convergence curves of the KD objective in the two KD processes.
  • Figure 5: Simulation results with JS divergence as the KD objective. (a), (b) and (c) plot the student's hidden states and the teacher's hidden states before and after the two KD processes. (d) shows the convergence curves of the KD objective in the two KD processes.
  • ...and 6 more figures