PaperlayerView exact revision on arXiv

Source-linked AI summary

Dual-Space Knowledge Distillation for Large Language Models

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

arXiv:2406.17328v3cs.CLcs.AI

TL;DR

Current white-box KD compares teacher and student distributions from different output spaces, limiting their similarity and requiring shared vocabularies. DSKD unifies those spaces and adds cross-model attention to align different vocabularies. On instruction-following benchmarks, it outperforms current white-box KD across distance functions and surpasses existing methods for different-vocabulary models.

  • Problem

    Current white-box KD compares distributions from different output spaces, which is associated with low teacher–student similarity, and requires the models to share a vocabulary.

    Source: §1 · p. 1

  • Method

    DSKD unifies teacher and student output spaces, with cross-model attention to align tokens for models with different vocabularies.

    Source: §1 · p. 2

  • Results

    On instruction-following benchmarks, DSKD significantly outperforms current white-box KD across distance functions for same-vocabulary models, while DSKD-CMA surpasses existing KD methods for different-vocabulary models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    With cross-model attention, the framework supports KD between LLMs regardless of their vocabularies and remains compatible with distribution-distance functions.

    Source: §1 · p. 2

    DSKD-CMA still lags slightly behind same-vocabulary DSKD in most cases.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Knowledge distillation offers a way to compress large language models by transferring a large teacher’s knowledge to a smaller student, but current white-box methods compare distributions from different output spaces. DSKD unifies those spaces, adds cross-model attention for vocabulary mismatch, and outperforms existing approaches on instruction-following benchmarks.

1 Introduction

  • KD transfers a large teacher’s knowledge and abilities to a smaller student, offering a way to compress LLMs amid rising deployment costs.The teacher supplies supervision signals during distillation.Source: §1 · p. 1
  • Different output spaces in current white-box KD limit teacher–student similarity at both representation and distribution levels, while the framework also requires shared vocabularies.Shared vocabularies are often not satisfied across current LLMs.Source: §1 · pp. 1–2
  • DSKD unifies the teacher’s and student’s output spaces by projecting their hidden states into each other’s representation spaces and using shared prediction heads.The framework remains compatible with distribution-distance functions such as KL and JS divergence.Source: §1 · p. 2
  • Cross-model attention automatically aligns tokens in differently tokenized sequences, enabling DSKD to support LLM pairs regardless of vocabulary.It also transforms their distributions to the same shape.Source: §1 · p. 2
  • DSKD significantly outperforms current white-box KD across distance functions for same-vocabulary models and surpasses existing KD methods for models with different vocabularies.The evaluation uses instruction-following benchmarks.Source: §1 · p. 2

2 Background and Preliminary Study

White-box KD trains a student to match a teacher’s token distributions, but separate output spaces may limit teacher–student similarity and prevent direct distillation across different vocabularies. Simulations show benefits from shared output spaces, while vocabulary differences create a separate alignment problem.

2.1 Current Framework for White-Box KD

  • White-box KD minimizes a distance between teacher and student token distributions, with temperature controlling their sharpness.The student distribution is pushed toward the teacher distribution at each token position.Source: §2.1 · p. 2

2.2.1 Low Teacher-Student Similarity

  • Using different prediction heads can limit similarity in both representations and distributions, according to the authors’ conjecture and simulation evidence.Different heads may yield dissimilar hidden states even when distributions are close, and make the theoretical minimum of the distribution distance difficult to reach.Source: §2.2.1 · p. 3
  • Shared-head KD made student hidden states more similar to the teacher and brought KL divergence closer to zero faster than different-head KD.The comparisons use simulation results with KL divergence; the figure shows hidden-state layouts and loss convergence.Source: §2.2.1 · p. 3, Figure 1 · p. 3

2.2.2 Dependency on the Same Vocabulary

  • Different vocabularies can tokenize the same text into different sequences, making teacher-token probabilities mismatched to student tokens and distribution dimensions incompatible.The authors state that the current white-box framework consequently fails for models with different vocabularies.Source: §2.2.2 · p. 3

3 Methodology

DSKD unifies teacher and student output spaces by projecting hidden states into the other model’s representation space and using a shared prediction head. Cross-model attention extends this approach to different vocabularies by learning token alignments.

3.1 Dual-Space Knowledge Distillation Framework

  • DSKD projects teacher and student hidden states into the other model’s representation space so shared prediction heads produce distributions in a unified output space.The framework combines KD losses in both spaces with the cross-entropy loss.Source: §3.1 · p. 4
  • In student space, a trainable linear projector maps teacher hidden states into the student representation space, and the student head produces the transformed teacher distribution.Because the projector is randomly initialized, that distribution is also trained to predict student-sequence target tokens with cross-entropy; the student-head gradient is stopped to avoid negative effects on the student.Source: §3.1 · p. 4
  • In teacher space, a separate projector maps student hidden states to the teacher dimension, whose prediction head produces student distributions for KD using KL divergence.The paper states that teacher-space KD fixes the distance function as KL divergence because it was found more appropriate there.Source: §3.1 · p. 4

3.2 Cross-Model Attention Mechanism

  • Cross-model attention learns token alignments from student embeddings and target tokens as queries, and teacher embeddings and hidden states as keys and values.The attention matrix weights teacher value vectors to produce aligned hidden states; the resulting representations support student- and teacher-space KD, with a mask used when teacher distributions are incorrect.Source: §3.2 · pp. 4–5

4 Experiments

The experiments compare DSKD with white-box and black-box KD across same- and different-vocabulary settings. DSKD outperforms current white-box KD across distance functions and, with cross-model attention, surpasses existing methods for different vocabularies.

4.1 Experimental Setup

  • The evaluation uses instruction-following data, with Dolly for KD training and additional datasets for testing.Dolly has about 11k training, 1k validation, and 500 test samples; Self-Instruct, VicunaEval, S-NI, and UnNI provide additional test sets.Source: §4.1 · p. 5

4.2 Baselines

  • Baselines include SeqKD, distance-function variants for same-vocabulary KD, and methods designed for different vocabularies.The comparisons include MinCE and ULD for different-vocabulary KD; on-policy KD methods are excluded because the study focuses on off-policy scenarios.Source: §4.2 · pp. 5–6

4.3 Results

  • All tested white-box KD methods outperform SeqKD in the same-vocabulary setting.The authors interpret this as evidence that token-level distributions transfer more knowledge than single target tokens.Source: §4.3 · p. 6
  • DSKD significantly outperforms current white-box KD across distance functions, while DSKD-CMA surpasses existing methods when model vocabularies differ.The results cover GPT2-120M and TinyLLaMA-1.1B students; the tables report Rouge-L scores across benchmarks.Source: §4.3 · pp. 6–7, Table 1 · p. 6, Table 2 · p. 7

5 Analysis

Analysis finds that unifying output spaces improves distillation results, and GPT-4 judgments favor DSKD in most tested cases. DSKD also yields more similar teacher and student representation structures than vanilla KD.

5.1 KD in Different Spaces vs. Unified Space

  • KD in the student space consistently outperforms KD in different spaces across distance functions, and adding teacher-space KD yields further gains.The student-space loss uses the same student distribution as current KD, but the teacher distribution shares the student’s output space.Source: §5.1 · p. 8

5.2 Evaluation via GPT-4

  • GPT-4 judges DSKD responses better than current white-box KD in most cases for both KL and reverse KL.The comparison uses responses to 100 randomly selected Dolly test instructions.Source: §5.2 · p. 8, Figure 2 · p. 8

5.3 Representation Similarity between the Teacher and the Student

  • DSKD yields lower distances between teacher and student representation structures than vanilla KD, indicating greater structural similarity.The study measures cosine-similarity and normalized-inner-product structures, using L1 distance averaged over 1,000 training samples.Source: §5.3 · p. 8, Figure 3 · p. 8

6 Related Work

Prior work applies white-box KD to language models and generation tasks, while LLM-specific methods have also emerged. This paper challenges the current framework and distinguishes DSKD from shared-head distillation by retaining the student’s original model size.

White-Box KD for Language Models

  • White-box KD for language models includes distribution-based and feature-based approaches, with applications in generation, translation, and summarization.The cited examples span pre-trained language models, machine translation, and text summarization.Source: §6 · p. 8
  • LLM-specific KD methods have been proposed, and this work challenges the current white-box framework with a different approach.The passage describes DSKD as a simple, more effective, and more general framework.Source: §6 · p. 9

KD with the Shared Prediction Head

  • Unlike SimKD, DSKD uses the teacher’s prediction head to transfer representation information without increasing the student’s original model size.SimKD equips the student with the teacher’s prediction head and therefore leaves a larger student with higher inference costs.Source: §6 · p. 9

7 Conclusion

The conclusion presents DSKD as a response to low teacher–student similarity and same-vocabulary requirements in current white-box KD. It unifies output spaces, adds cross-model attention for vocabulary mismatch, and reports gains across instruction-following benchmarks.

7 Conclusion

  • Current white-box KD is limited by low teacher–student similarity and a requirement that both models share a vocabulary.These limitations motivate the proposed framework.Source: §7 · p. 9
  • DSKD unifies teacher and student output spaces, while cross-model attention addresses vocabulary mismatch and supports distillation regardless of vocabulary.Source: §7 · p. 9
  • On instruction-following benchmarks, DSKD outperforms current white-box KD across distance functions and surpasses existing methods for different-vocabulary LLMs.Source: §7 · p. 9

Limitations

DSKD-CMA supports distillation between models with different vocabularies, but its performance usually remains slightly below DSKD with matching vocabularies.

Limitations

  • DSKD-CMA usually trails same-vocabulary DSKD slightly, a gap attributed to alignment errors between differently tokenized sequences.The authors describe cross-model attention as a relatively effective approach and suggest it may inspire future methods.Source: Limitations · p. 9

A.1 Simulation Results for Other Distance Functions

Across several alternative distance functions, the simulations report that conventional KD leaves a representation gap between student and teacher; the appendix also describes the simulation implementation and benchmark data.

A.1 Simulation Results for Other Distance Functions

  • Reverse KL, JS, skewed KL, skewed reverse KL, and adaptive KL simulations consistently show low student–teacher representation similarity after KD with different prediction heads.The results also show a gap from the minimum distribution distance; the authors suggest shared prediction heads may improve KD.Source: §A.1 · p. 12, Figure 4 · p. 12, Figure 5 · p. 12

A.2 Pseudo Code for Simulation Experiments

  • The simulation compares KD with different prediction heads against KD using the student's shared head, optimizing student models with SGD at learning rates in [1.0, 40.0] across distance functions.Source: §A.2 · pp. 13–14

B.1 Data

  • Evaluation datasets range from 80 Vicuna-Evaluation samples to 10,000 Unnatural Instructions samples, with three other test sets also included.Dolly contains 500 samples, Self-Instruction 242, and Super-Natural Instructions 1,694 samples with response lengths of at least 11.Source: §B.1 · p. 14

B.2 Training

The experiments specify training and evaluation procedures, select a temperature using validation performance, and compare KD variants across benchmarks. Student-space KD outperforms vanilla KD across all tested distance functions.

B.2 Training

  • The training section provides GPT2 and TinyLLaMA configurations and combines target-token cross-entropy with the KD loss for white-box methods.The experiments do not use the pre-training corpus during distillation.Source: §B.2 · p. 14, Table 4 · p. 14

B.3 Evaluation

  • Responses are randomly sampled using decoding temperature and top_p of 1.0, and Rouge-L scores are averaged across five random seeds.The seeds are 10, 20, 30, 40, and 50.Source: §B.3 · p. 14

B.4 Effect of Temperature for KD

  • Both KL and reverse KL perform best at temperature 2.0 on the validation set, so the experiments use 2.0 for all objectives.The reported search compares temperatures 1.0, 1.5, 2.0, 3.0, and 4.0.Source: §B.4 · p. 14, Figure 9 · p. 15

C Full Results

  • Across all distance functions, student-space KD outperforms vanilla KD, while teacher-space KD yields limited gains for some objectives.The authors select KL divergence for teacher-space KD because it performed relatively well there; that optimization does not directly affect the original student distribution.Source: C Full Results · p. 15

D Details and Full Results for GPT-4 Evaluation

GPT-4 pairwise evaluations favor DSKD over the current white-box KD framework across distance functions. The paper also details how it compares student and teacher representation structures.

D Details and Full Results for GPT-4 Evaluation

E Details of the Distance between Representation Structure

  • Representation-structure comparisons use cosine and normalized inner-product matrices, their L1 distances, and averages over 1,000 training samples.Smaller distances indicate more similar student and teacher representations; the comparison includes SFT, vanilla KD, and DSKD.Source: E Details of the Distance between… · p. 16
Loading 2406.17328v3…