Technical Summary: Rethinking the Role of Temperature in Large Language Model Distillation

Problem Statement

Knowledge distillation (KD) is a standard paradigm for compressing large language models (LLMs) by transferring knowledge from a teacher to a smaller student. A critical component of classical KD is the temperature parameter ($\tau$), which softens the teacher's probability distribution to reveal "dark knowledge" (informative signals beyond the top prediction). However, in recent LLM distillation literature, the temperature is frequently omitted or fixed at $\tau=1$.

This convention has led to a prevailing empirical conclusion that Reverse Kullback–Leibler (RKL) divergence is intrinsically superior to Forward KL (FKL) for LLM distillation. The authors argue that this conclusion is likely an artifact of evaluating these objectives under insufficiently softened teacher distributions ($\tau=1$). At $\tau=1$, teacher soft targets are sharp, suppressing the non-target probability mass that FKL relies on, potentially causing FKL to appear inferior not due to inherent limitations, but due to the evaluation setting. The paper seeks to determine if FKL has been systematically underestimated and how temperature fundamentally alters the comparison between FKL and RKL.

Methodology

The authors employ a combination of theoretical analysis and empirical evaluation to investigate the interaction between KD temperature and KL-based objectives.

Theoretical Analysis

The paper analyzes the gradients of FKL ($L_{FKL} = D_{KL}(p_\tau \| q_\tau)$) and RKL ($L_{RKL} = D_{KL}(q_\tau \| p_\tau)$) with respect to student logits.

• High-Temperature Regime: The authors prove that as $\tau \to \infty$, both FKL and RKL gradients converge to the same logit-matching behavior ($\nabla L \propto \frac{1}{V\tau^2}(z_s - z_t)$), making the objectives equivalent.
• Low-temperature Regime ($\tau \approx 1$): The objectives diverge significantly.
  • FKL: The gradient depends directly on the difference between the softened teacher distribution ($p_\tau$) and the student distribution ($q_\tau$). Increasing $\tau$ changes the supervision target itself, redistributing probability mass to non-dominant tokens and enriching the signal for FKL.
  • RKL: The gradient is driven by teacher-student logit gaps weighted by the student distribution ($q_\tau$). Here, temperature primarily acts to rescale the gradient and smooth the student-side weights, rather than fundamentally altering the target distribution in the same way as FKL.

Empirical Evaluation

Experiments were conducted on instruction-following benchmarks (Dolly Eval, Self-Instruct, Vicuna Eval, Super-NI, UnNI) using GPT-2 and OPT model families (e.g., GPT-2 XL $\to$ GPT-2 Base, OPT 6.7B $\to$ OPT 1.3B).

• Setup: Off-policy distillation where students are trained on fixed teacher-generated responses.
• Variables: The study compares standard objectives (FKL, RKL, Sym-KL, JS, SFKL, SRKL, AKL, AB, DRKL) at $\tau=1$ versus temperature-tuned variants (grid search over $\tau \in \{1.5, \dots, 4.0\}$).
• Metrics: ROUGE-L scores averaged over five random seeds.

Key Contributions

1. Asymmetric Effect of Temperature

The paper establishes that temperature affects FKL and RKL asymmetrically.

• FKL: Temperature substantially enriches the learning signal by exposing non-dominant token information from the teacher. It changes what FKL learns.
• RKL: Temperature mainly rescales the gradients and smooths weights without providing direct access to the softened teacher distribution as a probability target. It changes how RKL learns.

2. Reversal of Performance Rankings

The authors demonstrate that the superiority of RKL over FKL is highly temperature-dependent.

• At $\tau=1$, RKL consistently outperforms FKL, aligning with prior literature.
• At higher temperatures (typically $\tau \in [2.5, 3.5]$), FKL consistently surpasses RKL across all model scales and benchmarks. For instance, on GPT-2 Base, FKL improved from 19.41 to 20.97 ROUGE-L at $\tau=3.0$, while RKL remained nearly flat.

3. Broad Improvement of KL-Based Objectives

Temperature is not merely a correction for FKL but a general factor that improves a broad family of KL-based objectives (including Sym-KL, JS, SFKL, and AKL).

• Simple KL-based methods with appropriate temperature tuning achieve performance competitive with recent state-of-the-art (SOTA) methods like AB-KD and DRKL.
• For example, on GPT-2 Base, $JS+\tau$ achieved an average ROUGE-L of 25.63, surpassing AB (23.91) and DRKL (25.16).

Results

• Performance Gains: Temperature tuning yielded significant gains for FKL-based objectives (e.g., +5.83 for FKL on GPT-2 Base, +6.62 for AKL). Gains for RKL-based objectives were modest (e.g., +0.49 for RKL).
• Qualitative Improvements: Case studies on the Unnatural Instructions benchmark showed that temperature-tuned models better followed instructions regarding output format and entity extraction compared to their $\tau=1$ counterparts, reducing hallucinations and omissions.
• Ranking Shift: The introduction of temperature altered the relative ranking of distillation objectives, suggesting that previous comparisons conducted at $\tau=1$ may not reflect the true potential of FKL-based methods.

Significance and Claims

The paper claims that the KD temperature is not a minor implementation detail but a central, objective-dependent factor that must be considered for fair comparison in LLM distillation.

• Re-evaluating FKL: The authors argue that FKL has been systematically underestimated in the LLM distillation community due to the omission of temperature. When properly softened, FKL can outperform RKL.
• Simplicity vs. Complexity: The findings suggest that simple KL-based objectives, when paired with appropriate temperature scaling, can compete with complex, SOTA distillation frameworks, potentially simplifying the design space for future research.
• Fair Benchmarking: The paper asserts that future comparisons of distillation objectives must include temperature tuning to avoid biased conclusions regarding the intrinsic superiority of specific objectives like RKL.

The authors acknowledge limitations, noting that experiments were restricted to teacher models up to 7B parameters and excluded on-policy sampling or stronger student initialization, leaving these as avenues for future work.