PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models

Taiqiang Wu, Chaofan Tao, Jiahao Wang, Runming Yang, Zhe Zhao, Ngai Wong

arXiv:2404.02657v4cs.CLcs.AI

TL;DR

The paper questions whether RKL's presumed mode-seeking behavior makes it preferable to FKL for LLM knowledge distillation, since the usual seeking intuitions do not apply to this setting. It proposes AKL, which adaptively combines FKL and RKL; the divergences share an objective but emphasize different distribution regions early on, and AKL outperforms baselines while improving generated-response diversity and quality.

  • Problem

    The usual claims that FKL is mean-seeking and RKL is mode-seeking rely on a unimodal student and continuous distributions, assumptions that do not hold for LLM knowledge distillation.

    Source: §1 · p. 1

  • Method

    AKL adaptively allocates weights to combine FKL and RKL based on the teacher and student distributions.

    Source: §1 · p. 2

  • Results

    AKL outperforms baseline methods across popular benchmarks, and GPT-4 scores show improved diversity and quality of generated responses.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    At early epochs, FKL focuses on the distribution head and RKL on its tail, motivating adaptive combination when training is limited.

    Source: §1 · p. 2, §4.1 · p. 4

    Experiments cover GPT-2 and LLaMA 6.7B, but not larger models such as LLaMA-2 70B because of limited resources.

    Source: Limitation · p. 9

Abstract

from arXiv · show

Kullback-Leiber divergence has been widely used in Knowledge Distillation (KD) to compress Large Language Models (LLMs). Contrary to prior assertions that reverse Kullback-Leibler (RKL) divergence is mode-seeking and thus preferable over the mean-seeking forward Kullback-Leibler (FKL) divergence, this study empirically and theoretically demonstrates that neither mode-seeking nor mean-seeking properties manifest in KD for LLMs. Instead, RKL and FKL are found to share the same optimization objective and both converge after a sufficient number of epochs. However, due to practical constraints, LLMs are seldom trained for such an extensive number of epochs. Meanwhile, we further find that RKL focuses on the tail part of the distributions, while FKL focuses on the head part at the beginning epochs. Consequently, we propose a simple yet effective Adaptive Kullback-Leiber (AKL) divergence method, which adaptively allocates weights to combine FKL and RKL. Metric-based and GPT-4-based evaluations demonstrate that the proposed AKL outperforms the baselines across various tasks and improves the diversity and quality of generated responses. Codes are available at \href{https://github.com/wutaiqiang/LLM_KD_AKL}{github}.

1 Introduction

The paper revisits whether FKL or RKL is better suited to LLM knowledge distillation, arguing that their usual mean-seeking and mode-seeking intuitions do not hold in this setting. It proposes adaptively combining the divergences because they focus on different parts of the teacher distribution early in training, and reports improved benchmark, diversity, and quality results.

1 Introduction

  • Knowledge distillation compresses language models by training a compact student to mimic a teacher, addressing the deployment costs and latency associated with large models.The introduction frames compression as retaining performance while reducing model size.Source: §1 · p. 1
  • AKL outperforms baseline methods across popular benchmarks, and GPT-4 scores show improved diversity and quality of generated responses.The paper also reports human-evaluation results showing improved diversity and quality.Source: §1 · p. 2
  • Neither FKL's mean-seeking nor RKL's mode-seeking behavior holds for LLM knowledge distillation, and both converge to the same objective after sufficient epochs.The usual seeking-behavior claims rely on a unimodal student distribution and continuous teacher and student distributions.Source: §1 · pp. 1–2
  • At the beginning of training, FKL focuses on the teacher distribution's head while RKL focuses on its tail, motivating adaptive combination under limited epochs.AKL adaptively allocates weights to combine FKL and RKL based on the teacher and student distributions.Source: §1 · p. 2

2 Related Work

Prior work applies knowledge distillation to language-model compression through black-box and white-box approaches, including logit alignment with FKL or RKL. This paper revisits those divergences and distinguishes AKL's head- and tail-based adaptive weighting from prior combinations based on entropy or fixed weights.

2.1 KD for LLM Compression

  • Black-box KD uses teacher responses to improve student training data, whereas white-box KD aligns teacher and student logits during token generation.The paper places its work in white-box KD and revisits FKL and RKL before proposing AKL.Source: §2.1 · p. 2

2.2 Forward KL and Reverse KL

  • Prior KD work used FKL, combined FKL and RKL directly, or proposed BD-KD's weighted combination for online distillation.BD-KD selects weights of 1 or λ based on teacher and student entropy.Source: §2.2 · p. 2
  • Unlike BD-KD's entropy-based weighting, AKL assigns adaptive weights using the head and tail parts of the distributions.Source: §2.2 · p. 2

3 Preliminary and Rethinking

The section revisits the rationale for distinguishing FKL from RKL, then argues that their KD objectives coincide at convergence even though they emphasize different distribution regions early in optimization.

3.1 Background

  • Prior work motivates RKL over FKL by claiming mode-seeking versus mean-seeking behavior when fitting a Gaussian student to a multimodal teacher.Source: §3.1 · pp. 2–3, Figure 1 · p. 3

3.2 Rethinking FKL and RKL

  • Because LLM KD uses discrete distributions and students are not guaranteed to be Gaussian, the paper rejects the assumed mean/mode distinction and shows both divergences converge to teacher logits.The toy experiments use the same initial student distribution and show convergence after 200 epochs regardless of the teacher distribution's shape.Source: §3.2 · p. 3, Figure 2 · p. 4

KD for LLMs?

  • Early in optimization, FKL prioritizes high-probability head tokens while RKL prioritizes the low-probability tail, although both ultimately converge.On long-tail toy data, both converge after 250 epochs.Source: KD for LLMs? · p. 4

4 Methodology

Because practical LLM training limits the number of epochs, the section motivates balancing FKL and RKL's different early focus. AKL adaptively combines them using head- and tail-region gaps.

4.1 Motivation

  • Because LLM training rarely affords many epochs, FKL and RKL's early head/tail focus leaves logits poorly aligned, motivating a way to balance their losses.Source: §4.1 · p. 4, Figure 3 · p. 5

4.2 Adaptive Kullback-Leiber Divergence

  • AKL adaptively weights FKL and RKL according to which distribution region has the larger teacher–student gap, assigning more weight to FKL when the head gap is larger.Source: §4.2 · p. 5
  • AKL identifies head tokens by sorting teacher probabilities to meet threshold μ, then computes separate pointwise gaps for the head and tail.The pointwise-gap function is denoted ϵ(·).Source: §4.2 · p. 5

5 Experiments

Experiments compare distillation methods on Dolly, S-NI, and UnNI using Rouge-L, with GPT-2 and TinyLLaMA students. AKL outperforms the baselines without extra parameters, particularly on the larger UnNI dataset.

5.2 Main Results

  • In Rouge-L results, AKL outperforms all baselines without extra parameters, especially on larger UnNI; with 10,000 samples, it beats FKL+AKL for GPT-2 and TinyLLaMA (p < 0.001: 5e-9 and 9e-5, respectively).Source: Table 1 · p. 6, §5.2 · p. 6

5.1 Experimental Setup

  • The evaluation reports Rouge-L scores across five different seeds on Dolly (500 samples), S-NI (1,694 samples), and UnNI (10,000 samples), using GPT-2 1.5B/120M and LLaMA 6.7B/TinyLLaMA 1.1B teacher-student pairs.Training data follows an instruction-response dataset built from databricks-dolly-15k; GPT-2 is trained for 20 epochs.Source: §5.1 · pp. 5–6

5.2 Main Results

  • For GPT-2 on Dolly, FKL scores 23.46 versus 22.62 for RKL, despite prior claims that RKL is more suitable.The paper describes the gap between FKL and RKL as minor.Source: §5.2 · p. 6
  • Combining FKL and RKL improves performance over either alone, consistent with their focus on complementary head and tail regions.The authors state that combining them helps the student learn more information.Source: §5.2 · p. 6

6 Extensive Analysis

The analyses test AKL’s weight assignment, computational overhead, generated-response quality, and performance across tasks. AKL-r underperforms AKL, which adds modest compute cost; GPT-4 and human evaluations, task results, and case analysis support AKL’s effectiveness.

6.1 Ablation: Flipping the Loss Weights

  • On GPT-2, Figure 4 shows that after flipping the loss weights, AKL-r performs worse than AKL and FKL+RKL, supporting the effectiveness of AKL’s weight assignment.Source: §6.1 · p. 7

6.2 Computing Costs

  • AKL requires 2.8G more GPU memory and 4 minutes more training time than FKL+RKL.The authors describe the added memory and training time as slight and consider the overhead acceptable.Source: §6.2 · p. 7, Table 2 · p. 7

6.3 Evaluation via GPT-4

  • GPT-4 scores show AKL outperforming all baselines in diversity and quality, while diversity scores do not assess answer correctness.Figure 5 presents average diversity and quality scores for TinyLLaMA responses.Source: §6.3 · p. 7, Figure 5 · p. 8

6.4 User Study

  • Human evaluators found that AKL wins or ties baseline responses in more than 75% of cases for both diversity and quality.Five experts independently rated shuffled response pairs without knowing their source models.Source: §6.4 · pp. 7–8, Table 3 · p. 8

6.5 Performance under Various Tasks

  • Across Dolly subtasks, AKL combines FKL and RKL’s strengths and exceeds the teacher on Information Extraction and Creative Writing; it also outperforms baselines on Winogrande, OpenBookQA, BoolQ, and ARC.FKL performs well on Closed QA and Summary, while RKL performs well on Brain Storming.Source: §6.5 · p. 8

6.6 Case Study

  • Case analysis of TinyLLaMA responses supports the finding that AKL improves generated-response diversity and quality, consistent with GPT-4 scores and the user study.Source: §6.6 · p. 8

7 Conclusion

In LLM knowledge distillation, FKL and RKL do not exhibit their customary mean-seeking and mode-seeking behaviors; instead, they share an objective and converge with enough training. Because they focus on different distribution regions early on, AKL combines them adaptively, and evaluations report improved response diversity and quality.

7 Conclusion

  • In LLM distillation, FKL and RKL share the objective of matching teacher logits, and both converge after enough epochs.The conclusion says their mean-seeking and mode-seeking behaviors do not hold in this setting.Source: §7 · p. 8
  • Early in training, FKL focuses on the head and RKL on the tail, motivating AKL’s adaptive combination of the two divergences.Empirical results across various students demonstrate AKL’s effectiveness.Source: §7 · p. 8
  • GPT-4 scores and human evaluations indicate that AKL improves generated-response diversity and quality; scaling to larger models and more tasks remains future work.Source: §7 · p. 8

Limitation

Experiments did not extend to larger language models such as LLaMA-2 70B because of limited resources.

Limitation

  • Limited resources left experiments on larger models, including LLaMA-2 70B, for future work.The reported experiments used GPT-2 and LLaMA 6.7B.Source: Limitation · p. 9

Ethics Statement

The ethics statement notes inherited bias risks from generative LLMs; the appendix also gives an FKL convergence condition, evaluation details, and a comparison against additional baselines.

Ethics Statement

  • The ethics statement says generative LLM risks include gender and representation bias and biases toward blind and deaf people.It also encourages deployment of risk-free LLMs.Source: Ethics Statement · p. 9

A FKL and RKL converge

B Evaluation details

  • Evaluation details identify Dolly, S-NI, and UnNI datasets and report generation results across five random seeds.Source: B Evaluation details · p. 13

C More Baselines

D BERTscore Results

AKL consistently outperforms the baselines on TinyLLaMA BERTScore results, and tested threshold values 0.45 and 0.55 also outperform FKL+RKL. The section describes how GPT-4 scores responses for diversity and quality.

D BERTscore Results

  • 0.002, 0.028, and 0.00006 are the reported p-values for AKL outperforming FKL+RKL on BERTScore results.AKL consistently outperforms all baselines on TinyLLaMA generated responses.Source: D BERTscore Results · p. 13, Table 5 · p. 14

E Sensitivity of µ in AKL

F Scores from GPT-4

  • GPT-4 scores diversity and quality using generated responses from five seeds; due to cost, quality scoring uses Open QA and Closed QA responses.The quality scores are averaged across the five responses for each input.Source: F Scores from GPT-4 · pp. 13–14

G Results on more benchmarks

AKL outperforms the baselines on additional TinyLLaMA benchmarks. Across Dolly subtasks, FKL and RKL each perform well in different areas, while examples illustrate AKL outputs.

G Results on more benchmarks

H Performance on Subtasks

I Cases

  • Examples show AKL preserving teacher-model artistry, correcting grammar, returning a correct classification answer, and maintaining creative-writing style.The cited cases include the Miami-in-Florida answer and rich suggestions in the teacher's style.Source: I Cases · p. 15, Table 9 · p. 16, Table 10 · p. 17
Loading 2404.02657v4…