ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via $\alpha$-$\beta$-Divergence

We sincerely thank the reviewer for recognizing the theoretical foundations, clarity of contributions and experiments, and the improved performance demonstrated by our method. Our response follows:

Q1: I wonder whether the observation is consistent across different datasets. I hope the author can provide more clues on some other datasets, which could further strengthen the practical applicability of ABKD

A1: Thank you again for your insightful question. The table below outlines the hyperparam settings across datasets, which generally align with the theoretical results in Sec.3 (e.g., small α and large β for language modeling tasks and large α and small β for image classification tasks).

To further support our claim and validate the effects of our method, we made our best effort in the past few days to study OpenLLaMA-8B→3B distillation. For a thorough comparison, we also considered several baselines mentioned by other reviewers (i.e., AlphaNet, BDKD, Jensen's KL, and AKL) during the rebuttal period.

Based on the hyperparam tuning guidelines (App.D) derived from theoretical guidance, we simply set the hyperparams to α = 0.3 and β = 0.6 for all datasets (without further tuning).

The results below show that our method outperforms others by 0.65-3.26, especially excelling in Dolly and Unnatural.

We thank the reviewer for the thoughtful comments and suggestions. Our response follows:

Q1: The experiments should include a comparison of the performance of ABKD in these degenerate cases with the original FKLD and RKLD

A1: When ABKD degenerates to FKLD and RKLD, its performance matches the original FKLD and RKLD (Tab.3 in the main text). Thus, we don't separately consider these cases.

Q2: Theorem 3.2: Informal statement lacks rigor; formal proof needs verification

A2: Thank you for your question! The complete version of Thm.3.2 is in App.G.2. We will include them in the main text and improve readability upon acceptance.

Q3: This suggests task-specific tuning, challenging the "universal" claim. Tab.4 shows that ABKD requires specific hyperparam selection, limiting its scalability to additional datasets and architectures

A3: It's a pity that the term 'universal' is misunderstood. Actually, we aim to unify a series of existing divergences in KD, not provide a set of task-agnostic parameters. In fact, most related works require parameter tuning for different datasets and architectures. As per the 'no free lunch' theorem, universal parameters are impractical. Please see A4 of Reviewer 71pK for more details on 'universal'.

To further support our claim, we completed the distillation experiment from OpenLLaMA-8B to 3B within the limited time, including AKL mentioned by the reviewer, as well as BDKD, AlphaNet, and Jensen's KL mentioned by other reviewers. The results below show that our method outperforms others by 0.65-3.26, especially excelling in Dolly and Unnatural.

Q4: The Need for Ablation on Model Architectures.

A4: Our method applies to various architectures (e.g., 17 teacher-student pairs used in our experiments). It only needs adjusting the distillation objective. We also provide an ablation study of α-β divergence (Tab.3 in the main text), which shows using only α or β imposes unnecessary constraints on hardness- or confidence-concentration, leading to suboptimal solutions.

Q5: Essential References Not Discussed

A5: Thank you for your valuable insight! We will cite these works in the final version and discuss the differences between AKL and our work. We also conducted experiments comparing AKL with our method when distilling GPT-2 XL into GPT-2, as shown below. Our method outperforms AKL across datasets by 0.43-7.35.

Q6: Baseline Variance: Some baselines (e.g., LSD, TTM) underperform in Fig.5; unclear if hyperparams were optimized.

A6: All baseline results in Fig.5 are from the original papers. Missing values are obtained by rerunning their code with necessary hyperparam tuning and reporting the average of 3 random runs.

Q7: Cross-Architecture Experiment

A7: We conducted distillation from ResNet50 to VGG8 within the limited time. The results below show that our method can improve the performance of previous methods by 0.15-0.89.

Q8: How does ABKD perform with α/β outside [0,1] (e.g., α=1.5, β=-0.5)?

A8: This is an interesting question. We focus on balancing FKLD and RKLD, which correspond to extreme cases for α = 1, α = 0, β = 0, and β = 1. Thus, an intuitive method is to search for params between [0, 1]. Our experiments validate this. To address the reviewers' concern, we also tested α > 1 and β < 0 when distilling ResNet56 to ResNet20 within the limited time, as shown below:

An overly large α weakens the hardness-concentration and an overly small β weakens the confidence-concentration effect, both may degrade distillation performance.

Q9: Computational Cost Analysis

A9: We compared the training costs of different methods on language modeling, as shown below.

Our method takes a similar amount of time as Vanilla KD but is 1.15x to 5.80x faster than others due to its simplicity. It only modifies the optimization objective without adding extra cost, while others like GKD and DISTILLM require sampling student outputs during training.

We thank the reviewer for taking the time to provide their helpful and valuable feedback. Our response follows (please see https://anonymous.4open.science/r/ICML-rebuttal-experiments/results.md for all rebuttal experiments):

Q1: Essential References Not Discussed

A1: We will include the discussion of these baselines (the experiments are in Q5) in the new version:

• Jensen's KL is a popular method that combines FKLD and RKLD: $D_{JSD}(p|q)=1/2D_{KL}(p|m)+1/2D_{KL}(q|m)$, where $p$ is teacher distribution, $q$ is student distribution and $m=(p+q)/2$.

• BDKD adjusts the weights of FKLD and RKLD during training based on the entropy difference between $p$ and $q$.

• We also note AKL mentioned by the reviewer Reviewer Ty7Q, which adjusts the weights of FKLD and RKLD based on class differences between $p$ and $q$.

The problem with Jensen's KL is that when $p$ and $q$ are far apart, the loss becomes constant, causing vanishing gradients and hindering convergence. BDKD and AKL can be seen as variants of our baseline, WSD. However, they still tend to overemphasize small probability in $p$ and $q$, as shown in Sec.3.3.

Overall, a simple linear addition of FKLD and RKLD can't fully resolve their issues. This work aims to first theoretically analyze their limitations and then explore a more effective divergence to address them.

Q2: More Details on hyperparams α and β (Weakness #1)

A2: Thank you for your insightful question! We agree on the importance of clarifying hyperparam selection.

First, our method needs comparable or fewer hyperparams than prior works: AlphaNet has 3 (ICML 21), DISTILLM has 2 (ICML 24), and GKD has 2 (ICLR 24).

Second, while additional hyperparams increase search cost, our theoretical insights (Secs.3, 4) provide tuning guidelines for different tasks (App.D). These insights help eliminate less likely hyperparams, improving search efficiency. Empirical hyperparam selections (e.g., α=0.2, β=0.7 for NLP and α=0.9, β=0.2 for vision) validate these.

Overall, our search cost is not higher than prior works, and our hyperparams have clear theoretical and practical significance, helping one better understand the effectiveness of their distillation objective.

Q3: Training Cost and Convergence Behavior (Weakness #2)

A3: Thank you for your insightful question!

Training Cost: Tab.1 in the link shows our method requires a similar time to Vanilla KD, but is 1.15x to 5.80x faster than other SOTA methods, as it only modifies the distillation loss, while others need sampling student outputs during training, adding extra computational cost.

Convergence: Fig.1 in the link shows that our method outperforms others at all training stages (especially in the early stages), showing faster convergence.

Q4: Novelty issue (Weakness #3)

A4: Our main contribution lies in unifying FKLD, RKLD, and other potential variants from a novel theoretical perspective: the hardness- and confidence-concentration effects. This new theory helps explain why traditional divergences fail and our method succeeds:

• FKLD has overly weak hardness- and confidence-concentration effects. It fails to focus on the target class, leading to wrong predictions.
• RKLD has overly strong hardness- and confidence-concentration effects. It overemphasizes the target class while ignoring the broader knowledge from non-target classes.
• The weighted sum of FKLD and RKLD tends to overemphasize the minimal values in both teacher and student distributions.
• α-divergence imposes an unnecessary sum-to-one constraint on hardness- and confidence-concentration effects, potentially limiting model performance.
• α-β-divergence offers a flexible way to adjust hardness- and confidence-concentration separately, enabling smooth interpolation between FKLD and RKLD.

These insights are rarely explored in existing literature.

Q5: Performance Improvement & More Baselines (Weakness #4)

A5: Thank you for your valuable suggestion! We added more baselines (the suggested AlphaNet, BDKD, Jensen's KL, and AKL mentioned by other reviewers) with the limited time. Results are averaged over 5 seeds.

Tab.2 (see the link for the full version) shows our method outperforms others by 0.42-1.76.

To further validate our method, we made our best effort in the past few days to study OpenLLaMA-8B→3B distillation.

Tab.3 (see the link for the full version) shows it outperforms others by 0.65-3.26, especially excelling in Dolly and Unnatural.