PaperlayerView exact revision on arXiv

Source-linked AI summary

BiLD: Bi-directional Logits Difference Loss for Large Language Model Distillation

Minchong Li, Feng Zhou, Xiaohui Song

arXiv:2406.13555v3cs.CLcs.AI

TL;DR

LLM distillation must address large deployment costs, long-tail logit noise, and existing methods’ limited use of internal ranking information. BiLD filters long-tail noise and constructs bidirectional differences from top-k teacher and student logits to compute KL-based distillation losses. BiLD achieves superior distillation performance using only the top-8 logits compared with vanilla KL loss using full logits and other distillation methods.

  • Problem

    LLM distillation must address large deployment costs, long-tail logit noise, and existing methods’ limited use of internal ranking information.

    Source: §1 · pp. 1–2, §3.2 · p. 4

  • Method

    BiLD filters long-tail noise and constructs bidirectional differences from top-k teacher and student logits to compute KL-based distillation losses.

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

  • Results

    BiLD achieves superior distillation performance using only the top-8 logits compared with vanilla KL loss using full logits and other distillation methods.

    Source: §6 · p. 8

  • Takeaways & Limitations

    BiLD provides an effective alternative to full-logit vanilla KL that captures key teacher knowledge more efficiently within the evaluated settings.

    Source: §6 · pp. 8–9

    BiLD requires access to teacher logits and shared teacher-student vocabularies, while increasing the number of clipped logits rapidly raises time overhead.

    Source: Limitations · p. 9

Abstract

from arXiv · show

In recent years, large language models (LLMs) have shown exceptional capabilities across various natural language processing (NLP) tasks. However, such impressive performance often comes with the trade-off of an increased parameter size, posing significant challenges for widespread deployment. Knowledge distillation (KD) provides a solution by transferring knowledge from a large teacher model to a smaller student model. In this paper, we explore the task-specific distillation of LLMs at the logit level. Our investigation reveals that the logits of fine-tuned LLMs exhibit a more extreme long-tail distribution than those from vision models, with hidden "noise" in the long tail affecting distillation performance. Furthermore, existing logits distillation methods often struggle to effectively utilize the internal ranking information from the logits. To address these, we propose the Bi-directional Logits Difference (BiLD) loss. The BiLD loss filters out the long-tail noise by utilizing only top-$k$ teacher and student logits, and leverages the internal logits ranking information by constructing logits differences. To evaluate BiLD loss, we conduct comprehensive experiments on 13 datasets using two types of LLMs. Our results show that the BiLD loss, with only the top-8 logits, outperforms supervised fine-tuning (SFT), vanilla KL loss, and five other distillation methods from both NLP and CV fields.

1 Introduction

LLM distillation addresses deployment costs, but language-model logits contain long-tail noise and ranking information that existing methods do not fully exploit. BiLD uses top-k logits and bidirectional logit differences, achieving strong results across diverse distillation settings.

1 Introduction

  • Knowledge distillation transfers knowledge from a large teacher to a smaller student, supporting performance preservation and lightweight deployment.Source: §1 · p. 1
  • Logits distillation has been extensively studied for vision models but remains less explored for LLMs because their structures, data distributions, and output spaces differ.Source: §1 · p. 1
  • LLM logits exhibit a noticeable long-tail distribution containing substantial noise beyond a small amount of key knowledge.Source: §1 · p. 1
  • Existing logits distillation methods often underuse internal logit ranking, although ranking affects token selection in top-k and top-p generation.Source: §1 · p. 2
  • BiLD filters long-tail noise and leverages ranking by computing KL divergence over reconstructed pairwise logit differences from top-k teacher and student logits.Source: §1 · p. 2, Figure 1 · p. 2
  • BiLD uses only the top-8 logits to outperform SFT, vanilla KL loss, and five other distillation methods across 13 NLP datasets and two open-source LLMs.The evaluated models are BLOOM and Qwen1.5.Source: §1 · p. 2

2 Related Works

Prior logits-distillation research developed methods for vision models and emerging LLM-specific alternatives, but these approaches face transferability, optimization, complexity, or hyperparameter challenges. BiLD instead compares logits differences, focuses on key teacher knowledge, and can complement other distillation paradigms.

2.1 Logits Distillation

  • Vision distillation methods such as DKD, NKD, and NormKD are difficult to apply directly to LLMs because model structure, data, and output space differ.Source: §2.1 · pp. 2–3
  • LLM-specific methods trade off behavior or practicality: RKL can induce mode-seeking, DistiLLM adds hyperparameters, and SinKD is computationally demanding for larger models.Source: §2.1 · p. 3
  • BiLD computes divergence between logits differences rather than logits themselves, focusing on key teacher knowledge without excessive hyperparameters.Source: §2.1 · p. 3

2.2 Other Distillation Methods for LLMs

  • White-box distillation uses internal representations and often structural similarity, whereas black-box distillation commonly learns from teacher output text.Source: §2.2 · p. 3
  • BiLD is classified as black-box distillation, serves as an alternative to vanilla KL loss, and can also integrate with white-box distillation methods.Source: §2.2 · p. 3

3 Methods

The method first identifies sharper long-tail structure and underused ranking information in LLM logits, then introduces BiLD to distill selected logits through bidirectional differences.

3.1 Brief Review of Logits Distillation

  • Vanilla logits distillation converts teacher and student logits into temperature-normalized probabilities and aligns them with KL divergence.This provides the baseline formulation that BiLD modifies at the logits-difference level.Source: §3.1 · p. 3

3.2 The Characteristics of LLMs’ Logits

  • LLM logits are substantially sharper than image logits, with text-logit kurtosis reported as 2–3 orders of magnitude higher in a toy comparison.The comparison uses ResNet-101 and Qwen-4B logits and also examines top-k proportions.Source: §3.2 · pp. 3–4, Table 1 · p. 4

3.3 Bi-directional Logits Difference Loss

  • BiLD filters long-tail noise by selecting top-k teacher logits and corresponding student logits, then constructs bidirectional logits differences for task-specific LLM distillation.The loss combines teacher-led and student-led logits-difference components before computing the final optimization target.Source: §3.3 · p. 4, §3.3.1 · pp. 4–5

3.3.1 Formal Definition

  • BiLD’s teacher-led branch sorts the teacher’s top-k logits, selects student values at matching positions, and computes internal pairwise differences.The formal procedure then normalizes the differences into probabilities and evaluates the teacher-led loss.Source: §3.3.1 · p. 4
  • The student-led branch reverses the selection direction by choosing top-k student logits and extracting teacher logits at the corresponding positions.Its differences and probabilities are used to calculate the student-led logits-difference loss.Source: §3.3.1 · pp. 4–5
  • The BiLD calculation takes teacher and student logits, temperature T, and clipping parameter k as inputs, then outputs the combined BiLD loss.Algorithm 1 summarizes selection, difference construction, normalization, and the two component losses.Source: §3.3.1 · p. 5

3.3.2 Explanation about the Utilization of Logits Ranking

  • The sign of the student-corresponding logits difference indicates whether teacher–student value-pair rankings agree, constraining the student toward the teacher’s ranking order.This ranking signal supplements matching the teacher’s selected logit values.Source: §3.3.2 · p. 5

4 Experiments

The experiments evaluate BiLD across 13 NLP datasets, multiple LLM families, and seven comparison baselines using shared training and distillation settings.

4.1 Datasets

  • The evaluation covers 13 datasets: eight from SuperGLUE and five additional model-compression datasets, evaluated separately after collective training to reduce overfitting.The datasets vary substantially in size, so SFT and distillation use them collectively before per-dataset evaluation.Source: §4.1 · p. 5

4.2 Baselines

  • BiLD is compared against SFT, vanilla KL, top-k KL, three vision-model distillation methods, and reverse KL used in MiniLLM.The vision-model baselines are DKD, NKD, and NormKD.Source: §4.2 · p. 5

4.3 Implementation Details

  • Experiments use BLOOM-7B and Qwen-4B teachers with BLOOM-3B, BLOOM-1B, Qwen-1.8B, and Qwen-0.5B students.These models provide two LLM families and multiple teacher–student size pairings.Source: §4.3 · pp. 5–6
  • Training uses three SFT epochs and eight distillation epochs with batch size 64, micro batch size 2, cosine scheduling, and temperature T=3.SFT uses cross-entropy, while distillation applies the tested loss to the output portion and vanilla KL to the instruction portion.Source: §4.3 · p. 6
  • The experiments use k=1024 for top-k KL and k=8 for BiLD, with DeepSpeed, gradient checkpointing, and BFLOAT16 on eight NVIDIA A100 GPUs.Evaluation can be performed with a single NVIDIA A100 GPU using vLLM.Source: §4.3 · pp. 6–7

5 Results and Analysis

Across 13 datasets and four distillation settings, BiLD achieves the strongest average accuracy while filtering long-tail logits and improving alignment with teachers’ important token rankings. Analyses show that temperature and top-k choices materially affect performance, motivating T=3 and k=8.

5.1 Main Results

  • BiLD achieves the highest average accuracy across all four distillation settings, outperforming SFT, vanilla KL, and five other methods.It surpasses vanilla KL by 3.52%, 1.09%, and 1.10% in three settings, while retaining a 0.64% advantage where vanilla KL nearly matches the teacher.Source: §5.1 · p. 7

5.2 Analysis of the Effectiveness of Clipping Logits

  • Clipping logits to the top-k improves KL performance in three distillation scenarios, indicating that filtering long-tail noise can enhance distillation.The top-1024 logits cover over 99% of teacher probability for both Qwen-4B and BLOOM-7B.Source: §5.2 · p. 7

5.3 Analysis of Performance at the Logit Level

  • BiLD notably improves overlap@8 while maintaining competitive overlap@1, indicating better imitation of teachers’ important logit behavior.Overlap@1 measures matching highest-logit tokens, whereas overlap@k for k>1 measures overlap among important teacher and student tokens.Source: §5.3 · pp. 7–8, Table 3 · p. 7

5.4.1 Impact of Temperature

  • Lower temperatures significantly degrade BiLD, while T=3 yields the best performance in the temperature ablation.The temperature sweep varies T over {0.1, 0.5, 1, 3, 5, 8, 10} with other settings fixed.Source: §5.4.1 · p. 8, Figure 2 · p. 8

5.4.2 Impact of the k Value in BiLD Loss

  • Small k values produce poor performance, while larger k improves and then stabilizes accuracy and overlap at increased computational cost.The study selects k=8 as the computation–performance trade-off for BiLD.Source: §5.4.2 · p. 8, Figure 3 · p. 8, Table 4 · p. 8

6 Conclusion

BiLD is a logits-distillation objective that filters long-tail noise and uses internal ranking information. Experiments across diverse datasets and architectures show effective knowledge transfer using only the top-8 logits.

6 Conclusion

  • BiLD achieves superior distillation performance using only the top-8 logits compared with vanilla KL loss using full logits and other methods.The loss filters long-tail noise and leverages internal logits ranking information.Source: §6 · p. 8
  • Experiments across diverse datasets and model architectures confirm BiLD’s effectiveness and its ability to capture key teacher knowledge more efficiently.Source: §6 · p. 9

Limitations

BiLD requires teacher-logit access and a shared vocabulary, limiting which teacher–student pairs it can support. Its computational overhead also increases rapidly with larger top-k selections, while clipping can discard long-tail knowledge.

Limitations

  • BiLD cannot use teachers such as GPT-4 and Gemini when they provide only text or incomplete logits, and it requires shared teacher–student vocabularies.Source: Limitations · p. 9
  • Increasing the number of clipped logits rapidly raises BiLD’s time overhead, making larger top-k settings a practical concern.Using only the top-8 logits performs better than vanilla KL loss, but increasing the clipped-logit count escalates overhead.Source: Limitations · p. 9
  • Directly clipping the long-tail logits improves performance but inevitably loses knowledge contained in that distribution.Better ways to use long-tail knowledge remain an open research direction.Source: Limitations · p. 9

A Calculation Efficiency of BiLD

For distillation from Qwen-4B to 0.5B, BiLD improves performance over vanilla KL with an acceptable training-time increase. Its main computational bottleneck is constructing internal pairwise logit differences.

A Calculation Efficiency of BiLD

B Details about Dataset Sizes

The datasets used in the experiments vary substantially in size. CB, COPA, and WSC are about three orders of magnitude smaller than ReCoRD.

B Details about Dataset Sizes

C Toy Experiment to Compare Vision Model and LLMs’ Logits

The toy experiment uses five images and five instruction sets to compare vision-model and LLM logits.

D Templates

The datasets use task-specific templates spanning multiple-choice, entailment, effect-prediction, and text-completion formats.

  • The traffic-law template asks whether turning right on red is legal in California using a true-or-false choice.Source: D Templates · p. 13
  • The text-completion template asks readers to select the most appropriate continuation for a procedural passage.The example concerns cleaning a plastic retainer and offers several candidate steps.Source: D Templates · p. 14
  • Several templates present a passage or premise with answer choices for classification or reasoning.Examples include funding questions, flooding solutions, fault-measurement tools, light reflection, traffic-law questions, entailment, and possible effects.Source: D Templates · pp. 13–14
  • The dataset templates are provided in Table 6.Source: D Templates · p. 13, Table 6 · p. 15

Related papers

Loading 2406.13555v3…