Source-linked AI summary
Rethinking Reverse KL as Adaptive Entropy Distillation
Shizhen Li, Zhiyu Shen, Yuyin Lu, Yunhe Pang, Jielin Song, Yanghui Rao, Fu Lee Wang
TL;DR
Existing on-policy distillation methods struggle to balance probability coverage and concentration while accounting for token-level teacher confidence. This paper decomposes RKL into an RKL-only adaptive objective, proposes AED using teacher entropy to calibrate imitation strength, and reports strong task performance and distributional alignment.
Problem
Existing on-policy distillation methods use fixed or FKL–RKL balancing strategies that do not account for varying token-level teacher confidence.
Method
AED reweights RKL’s internal concentration and coverage components and uses teacher entropy to adaptively calibrate token-level imitation strength.
Results
AED achieves strong performance in task evaluation, distributional alignment, and entropy alignment across instruction-following and mathematical reasoning benchmarks.
Takeaways & Limitations
The results support RKL-only adaptive entropy distillation as a practical approach for improving task performance and teacher–student alignment.
Takeaways & Limitations
AED requires open-source teacher and student models with consistent vocabularies, and experiments did not scale to ultra-large models because of computational constraints.
Abstract
from arXiv · showhide
Knowledge distillation (KD) is widely used to transfer the capabilities of large language models (LLMs) to smaller students, but existing objectives often struggle to balance faithful imitation and robust generation. In particular, existing methods mainly combine FKL and RKL, overlooking that RKL itself provides a mechanism for adjusting the student's imitation strength. Motivated by this, we revisit on-policy Reverse Kullback-Leibler (RKL) distillation and decompose its objective into a teacher-fitting term and a student-entropy term, without introducing an explicit FKL branch. We show theoretically that the token-level optimal student distribution corresponds to a tempered variant of the teacher distribution, where the adaptive weight controls the trade-off between mode-seeking and uncertainty preservation. Guided by this insight, we propose \textbf{Adaptive Entropy Distillation (AED)}, which uses the teacher's entropy to dynamically calibrate token-level imitation strength. Experiments on instruction-following and mathematical reasoning benchmarks demonstrate that AED achieves superior overall performance and generally improves teacher--student distributional and entropy alignment.
1 Introduction
The introduction identifies fixed global FKL/RKL weighting as unable to reflect token-level uncertainty and argues that RKL itself can adapt imitation strength. It presents AED, an RKL-only method that reweights RKL components using teacher entropy and improves performance and distributional alignment.
- Motivation: Fixed-weighting distillation uniformly scales objectives across sequences, failing to account for varying token-level uncertainty.Existing approaches use static FKL/RKL interpolations or variants such as Skewed RKL.
- Core insight: RKL’s internal structure provides a mechanism for controlling the student’s imitation strength without explicitly combining FKL and RKL.The introduction reframes RKL as containing complementary forces for concentration and coverage.
- Theory: The reweighted RKL objective has an analytical optimum corresponding to a tempered teacher distribution.This characterization explains how reweighting controls the student’s imitation strength.
- Method: AED adaptively maps teacher entropy to token-level weights, calibrating imitation strength within an RKL-only objective.The framework is theory-driven and uncertainty-aware, and is designed to balance RKL’s internal components.
- Results: Experiments show that AED improves task performance and teacher–student distributional alignment.The introduction attributes these gains to reweighting RKL’s intrinsic components and calibrating imitation strength with teacher entropy.
2 Related Work
Related work has progressed from off-policy knowledge distillation to on-policy methods using student-generated samples and teacher supervision. Studies have also explored alternative divergences and adaptive weighting, while this work analyzes RKL’s distillation behavior beyond treating it as fixed.
- Traditional knowledge distillation: Traditional KD typically uses static teacher-generated targets, including soft logits, intermediate representations, rationales, and synthetic instruction-following data.For LLMs, prior work applies KD to chain-of-thought reasoning and instruction-following behaviors.
- On-policy distillation: Recent on-policy distillation methods let students sample from their own distributions while receiving teacher supervision.MiniLLM uses token-level RKL corrective signals, whereas GKD generalizes on-policy distillation with alternative divergence objectives.
- Alternative divergence objectives: Prior work explores Skewed RKL, JS divergence, SKL, and general f-divergences to balance RKL’s mode seeking with broader distributional coverage.These objectives appear in DistilLLM, GKD, and related on-policy distillation studies.
- Adaptive divergence balancing: Adaptive approaches balance divergence behaviors dynamically, including AKL’s head-versus-tail adaptation of FKL and RKL fitting.The paper distinguishes its approach by decomposing and analyzing RKL rather than treating it as a fixed objective.
3 Preliminaries and Analysis
The section formalizes on-policy token-level distillation and analyzes how FKL and RKL shape student training. It shows that RKL combines teacher fitting with student entropy, motivating direct control of their trade-off.
- Problem setup: On-policy distillation samples responses from the student and evaluates teacher and student vocabulary distributions at each prefix.Sequence-level objectives decompose into token-level terms along student-sampled trajectories.
- RKL analysis: RKL decomposes into teacher-derived cross-entropy and negative student entropy, balancing mode seeking against probability spreading.The cross-entropy penalizes mass on teacher-unlikely tokens, while negative entropy penalizes overconcentration.
- FKL analysis: On-policy FKL uses teacher-weighted cross-entropy to broaden student coverage of the teacher distribution, but its entropy term is constant and supplies no student gradient.This support-covering pressure may produce mean-seeking behavior when the student is less expressive than the teacher.
- RKL analysis: Standard RKL assigns equal relative weights to its teacher-fitting and student-entropy terms, suggesting that concentration-versus-coverage trade-offs can be controlled through adaptive reweighting.This observation motivates the formulation introduced in the next section.
4 Method
AED is an on-policy distillation framework that decomposes RKL into teacher fitting and student entropy, yielding a tempered teacher target whose sharpness is controlled by adaptive imitation strength. It calibrates this strength from teacher entropy, strengthening guidance for low-entropy tokens and broadening coverage for high-entropy tokens.
- Theoretical characterization: At a fixed student-generated prefix, minimizing the adaptive token-level loss yields a tempered variant of the teacher distribution.The teacher distribution and adaptive weight are treated as fixed in the local analysis.
- Theoretical characterization: Minimizing the adaptive loss is equivalent to minimizing RKL divergence to the corresponding tempered teacher distribution.The equivalence holds because 1 − λt is positive and the normalization term is independent of qt.
- Adaptive target: The adaptive weight λt controls inverse temperature αt: larger values sharpen imitation toward primary teacher modes, while smaller values encourage broader probability coverage.AED therefore implements a controllable mode-seeking paradigm within the RKL framework.
- Framework: AED decomposes RKL into teacher fitting and student-entropy terms, with the tempered teacher distribution emerging as an equivalent characterization rather than a separately imposed temperature objective.AED directly optimizes the adaptive loss in Equation 5.
- Entropy-driven calibration: AED maps teacher token entropy to λt ∈ [0, 1] using a bounded, monotonically decreasing function, so uncertainty systematically reduces imitation strength.The mapping is instantiated with normalized divergence between the teacher distribution and the uniform distribution, avoiding additional mapping hyperparameters.
- Entropy-driven calibration: Low-entropy teacher tokens produce λt near 1 and stronger teacher guidance, whereas high-entropy tokens produce smaller λt and increase the student-entropy contribution.This entropy-driven coupling adapts imitation strength to the teacher’s local distributional characteristics.
5 Experiments
Experiments evaluate AED across instruction-following and mathematical reasoning tasks, comparing it with RKL and recent adaptive distillation baselines. AED generally delivers stronger performance, improves teacher–student alignment, and benefits from adaptive weighting over fixed or heuristic alternatives.
- Evaluation Setup: AED is evaluated on five instruction-following benchmarks using ROUGE-L and LLM-as-a-Judge pairwise preference evaluation.The evaluated benchmarks are Dolly, SelfInst, Vicuna, S-NI, and UnNI.
- Overall Performance: Across model architectures and parameter scales, AED achieves favorable overall performance and the highest average scores.The experiments compare three teacher–student configurations and include standard RKL, AKL, ToDi, and EOPD baselines.
- LLM-as-a-Judge Evaluation: AED achieves higher win rates than loss rates against RKL, AKL, ToDi, and EOPD in LLM-as-a-Judge comparisons.The evaluation uses GPT-5.4 as judge with the TinyLLaMA-1.1B student on S-NI and UnNI subsets.
- Ablation Study: Standard AED significantly outperforms fixed-weight RKL and the adaptive FKL+RKL variant across most benchmarks.The fixed λ = 0.5 baseline misses token-level trade-offs, while FKL can pressure capacity-limited students to cover teacher support rather than concentrate on high-confidence modes.
- Alignment and Generalization: AED achieves the lowest RKL and MAEG values in GPT-2 and LLaMA2–TinyLLaMA settings, competitive alignment under Qwen3, and the best overall mathematical reasoning performance.On MATH-500, AMC 2023, and AIME 2024, AED also achieves higher Pass@1 and consistently higher Pass@k under the same sampling budget.
6 Conclusion
The paper decomposes RKL into a tempered framework and introduces AED, which adapts token-level distillation intensity to teacher uncertainty. Experiments show strong task, distributional, and entropy alignment performance.
- 6 Conclusion: AED adaptively adjusts token-level distillation intensity according to teacher uncertainty.This method follows the decomposition of the RKL objective and its equivalence to a tempered RKL framework.
- 6 Conclusion: Experiments demonstrate strong performance in task evaluation, distributional alignment, and entropy alignment.
Limitations · A Analytic Optimum of the Token-Level Adaptive Objective · A.1 Lagrangian Optimization for the Adaptive Objective
The section identifies AED’s vocabulary and scaling limitations, then derives the token-level analytic optimum by constrained Lagrangian optimization. The resulting adaptive objective is mathematically equivalent to reverse RKL against a temperature-scaled teacher distribution.
- Limitations: AED requires teacher–student vocabulary consistency for token-level alignment, while extension to black-box APIs and heterogeneous vocabularies remains open.Experiments also were not scaled to ultra-large models because of computational resource constraints.
- Limitations: The theoretical formulations are scale-agnostic, although experiments did not include ultra-large models.This limitation reflects computational resource constraints rather than a stated theoretical restriction.
- A Analytic Optimum of the Token-Level Adaptive Objective: The token-level adaptive loss has an analytic optimum that is mathematically equivalent to minimizing reverse RKL against a temperature-scaled teacher distribution.This establishes the optimization interpretation of the adaptive objective.
- A.1 Lagrangian Optimization for the Adaptive Objective: For a fixed prefix y<t, pt(v) and qt(v) denote teacher and student probabilities over vocabulary V at token step t.The token-level expected loss is defined over these distributions.
- A.1 Lagrangian Optimization for the Adaptive Objective: The optimal student distribution q∗t is obtained by minimizing the token-level loss subject to the probability-simplex constraint Σ_v qt(v) = 1.The optimization is performed for each fixed prefix and token step.
- A.1 Lagrangian Optimization for the Adaptive Objective: A Lagrangian is constructed for the constrained objective, and setting the partial derivative with respect to qt(v) to zero characterizes the optimum.The derivation proceeds through the stationarity condition of the constrained problem.
- A.1 Lagrangian Optimization for the Adaptive Objective: Solving the stationarity equation for log qt(v) introduces a constant C independent of v.The supplied derivation specifies C = −1 − η 1−λt.
- A.1 Lagrangian Optimization for the Adaptive Objective: The analytic optimum explicitly takes the form of a scaled student distribution.The scaling form follows after solving for log qt(v) and identifying the v-independent constant.
A.2 Mathematical Equivalence to Tempered RKL · B Implementation Details
The adaptive token-level loss is mathematically equivalent to reverse KL against a temperature-scaled teacher distribution. Teacher entropy dynamically adjusts mode-seeking strength, smoothing uncertain targets while preserving confident modes.
- A.2 Mathematical Equivalence to Tempered RKL: The optimality condition shows that the adaptive objective minimizes divergence against a normalized, temperature-scaled teacher distribution.This establishes the tempered teacher as the adaptive objective’s implicit target.
- A.2 Mathematical Equivalence to Tempered RKL: Rewriting the original loss with the tempered teacher distribution and its logarithm exposes the normalization term required for equivalence.The derivation uses the logarithmic relation between the original teacher probabilities and the tempered distribution.
- A.2 Mathematical Equivalence to Tempered RKL: Because the student distribution sums to one, the first summation term becomes the standard Kullback-Leibler divergence.This normalization property is the key algebraic step before simplifying the expected objective.
- A.2 Mathematical Equivalence to Tempered RKL: The expected adaptive loss simplifies to an objective whose minimization is equivalent to minimizing RKL between the student and tempered teacher distributions.Thus, the method can be interpreted as tempered RKL rather than as an explicit combination with an FKL branch.
- A.2 Mathematical Equivalence to Tempered RKL: Compared with direct RKL targeting, the objective provides a controllable mode-seeking paradigm through the adaptive inverse temperature αt.The inverse temperature is adjusted in response to teacher entropy.
- A.2 Mathematical Equivalence to Tempered RKL: High-uncertainty regions receive safely smoothed targets, whereas low-uncertainty regions preserve confident mode-seeking behavior.This behavior follows from dynamically adjusting αt according to teacher entropy.
B.1 Training details
Training uses the databricks/dolly-15k instruction-following dataset on two NVIDIA RTX A100 (80GB) GPUs, with model-specific compute costs and configurations reported for GPT-2, TinyLLaMA, and Qwen3.
- Instruction Following: The instruction-following experiments train on databricks/dolly-15k, comprising 11K training, 1K validation, and 500 test samples.
- Instruction Following: Training runs use two NVIDIA RTX A100 (80GB) GPUs.
- Instruction Following: 3 GPU hours, 30 GPU hours, and 21.5 GPU hours are the approximate training costs for GPT-2, TinyLLaMA, and Qwen3, respectively.
- Instruction Following: Detailed training configurations for the three model architectures are summarized in Table 4.
B.2 Evaluation details
Evaluation uses fixed sampling settings and final checkpoints across instruction-following and mathematical-reasoning benchmarks. Instruction-following results average ROUGE-L over five seeds, while mathematical reasoning reports Avg@8 and Pass@8 from eight sampled responses per problem.
- Instruction Following: Instruction-following evaluation samples responses at temperature 1.0, top-k 0, and top-p 1.0.Full test sets are used for DollyEval, SelfInst, Vicuna, and S-NI; UnNI uses its first 10,000 examples.
- Instruction Following: Five random seeds are used to account for generation randomness, with mean ROUGE-L and standard deviation reported across runs.The seed set is {10, 20, 30, 40, 50}, and the final training checkpoint is evaluated.
- Mathematical Reasoning: Mathematical-reasoning evaluation generates eight responses per problem with maximum length 2,048 tokens, temperature 1.0, top-k 0, and top-p 0.8.Results report Avg@8 as average accuracy and Pass@8 as the proportion of problems with at least one correct response, using the final checkpoint.
C Distributional and Entropy Alignment in Mathematical Reasoning
On AMC23 mathematical reasoning distillation, the study evaluates token-level RKL and entropy alignment using MAEG. AED achieves the best alignment under evaluation settings otherwise following Section 5.4.
- Evaluation setup: The evaluation measures token-level RKL and entropy alignment on AMC23 using the MAEG.Except for the dataset, the evaluation settings follow those in Section 5.4.
- Results: AED obtains the best alignment among the evaluated methods.This result is reported for the AMC23 mathematical reasoning distillation evaluation.