ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via $\alpha$-$\beta$-Divergence
Authors: Guanghui Wang
Zhiyong Yang
Zitai Wang
Shi Wang
Qianqian Xu
Qingming Huang
Venue: ICML 2025 oral (2025)
Keywords: Knowledge Distillation
α-β divergence
Primary area: general_machine_learning

Abstract: Knowledge Distillation (KD) transfers knowledge from a large teacher model to a smaller student model by minimizing the divergence between their output distributions, typically using forward Kullback-Leibler divergence (FKLD) or reverse KLD (RKLD). It has become an effective training paradigm due to the broader supervision information provided by the teacher distribution compared to one-hot labels. We identify that the core challenge in KD lies in balancing two mode-concentration effects: the \textbf{\textit{Hardness-Concentration}} effect, which refers to focusing on modes with large errors, and the \textbf{\textit{Confidence-Concentration}} effect, which refers to focusing on modes with high student confidence. Through an analysis of how probabilities are reassigned during gradient updates, we observe that these two effects are entangled in FKLD and RKLD, but in extreme forms. Specifically, both are too weak in FKLD, causing the student to fail to concentrate on the target class. In contrast, both are too strong in RKLD, causing the student to overly emphasize the target class while ignoring the broader distributional information from the teacher. To address this imbalance, we propose ABKD, a generic framework with $\alpha$-$\beta$-divergence. Our theoretical results show that ABKD offers a smooth interpolation between FKLD and RKLD, achieving a better trade-off between these effects. Extensive experiments on 17 language/vision datasets with 12 teacher-student settings confirm its efficacy.

ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via a-S-Divergence
Guanghui Wang'
Zhiyong Yang'
Zitai Wang?
Shi Wang?
Qianqian Xu
Qingming Huang 132
Abstract
its efficacy. The code is available at https:
/ //github. com/ gww.//gwwwgg//////ae abkd.
Knowledge Distillation (KD) transfers knowl- edge from a large teacher model to a smaller student model by minimizing the divergence be- tween their output distributions, typically using forward Kullback-Leibler divergence (FKLD) or reverse KLD (RKLD). It has become an effec- tive training paradigm due to the broader super- vision information provided by the teacher distri- bution compared to one-hot labels. We identify that the core challenge in KD lies in balancing two coneetonncnttttion effects: the Hardness- Concentration effect, which refers to focusing on modes with large errors, and the Confidence- Concentration effect, which refers to focusing on modes with high student confidence. Through an analysis of how probabilities are reassigned dur- ing gradient updates, we observe that these two effects are entangled in FKLD and RKLD, but in extreme forms. Specifically, both are too weak FKLD, causing the student to fail concentrate the target class. In contrast, both too strong in RKLD, causing the student to overly empha- size the target class while ignoring the broader distributional information from the teacher. To address this imbalance, we propose ABKD, a generic framework with a-3 -i-deivegeence. Our the- oretical results show that ABKD offers smooth interpolation between FKLD and RKLD, achiev- ing effective trade-off between these effects. Extensive experiments on 17 language/vision datasets with 12 teacher-student settings confirm
1. Introduction
Knowledge Distillation (KD) (Hinton, 2015) is widely- adopted technique for transferring knowledge from large models (teachers) to smaller models (students). In this setup, the student model, with a predictive distribution qo, learns to mimic the predictive distribution p of the teacher model. This imitation is typically achieved by minimizing a prede- fined divergence ID between the teacher distribution p and the student distribution qo: fKD - D(pllg)). This way, KD allows the student to leverage richer soft label information from p compared to one-hot labels, often leading to better performance than traditional supervised fine-tuning. This has been shown in tasks like image classification (Dosovit- skiy, 2020; Radford et al., 2021; Yang et al., 2023b; Wang et al., 2022b) and text generation (Vaswani, 2017; Touvron et al., 2023a).
A key step in KD is to choose proper divergence D for dis- tribution matching. One popular choice in previous works (Cho & Hariharan, 2019; Mirzadeh et al., 2020; Zhou et al., 2021; Zhao et al., 2022; Jin al., 2023; Sun et al., 2024; Zheng & Yang, 2024) is the forward Kullback-Leibler diver- gence (FKLD). However, FKLD's asymmetry often results in student distribution qe that is overly smooth, spreading across the entire support of p. To address this, recent studies (Lee et al., 2023; Gu et al., 2024a; Kim et al., 2024; Gu et al., 2024b) have explored the reverse KI.D (RKLD), which al- lows qo to focus on a few prominent modes of p. Despite the effectiveness,, empirical results (Wen al., 2023; Wu et al., 2024; Ko et al., 2024) suggest that RKLD often yields suboptimal performance across a range of tasks. What is worse, there is systematic approach to identify the essen- tial issues hidden behind, which hinders the development of a more generic and effective KD framework. To get out of this dilemma, we first pose the following question:
'Scooo of Computer Science and Technology, University of Chinese Academy of Sciences, Beijing, China Key Laboratory of Intelligent Information Processing, Institute of Computing Tech- nology, Chinese Academy of Sciences, Beijing. China Key Labo- ratory of Big Data Mining and Knowledge Management (BDKM), University of Chinese Academy of Sciences, Beijing, China. Cor- respondence to: Zhiyong Yang <<ynnzzhyong221 (uucas.ac.cn>> Qingming Huang <qmhuang uc.>..
What underlying factors contribute to the suboptimal performance of FKLD and RKLD?
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
To answer this, we analyze how different divergence func-

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via
(a)
Attention of enor wxte
(b)
a-B. -ivviieeeeecce
probability mass Tarpet Distribution
(d)
probability mass
(e)
Target Detributioo RKLD
a-
28,0
FKLD
RKLD
24.0
ener (hardhees)
(c)
output
output
20.0
Acannion candderce
probability mass
(f)
probability mass
18.0
o-divergence
Targer DSUu busion
Targat Datr bution e-B divergence
02
beta
0.4
0.6
1.0
1.0
Manbeess
08
alpha
04
02
0.0
Foutput
output
Figure 1. (a) Illustration of the unified search space for our proposed ABKD, where height (color) represents performance (t). The FKLD
and RKLD are special cases of ABKD when selecting (a = 1, = 0) and (a = 0, = 1) respectively. The a-divergence can
only search along the submanifold a + 3 = 1 in the ABKD space. (b)-(c) illustrate how adjusting a and 3 affects narddnes--ooncetration and cewpoddence--connntrtttiio (d)-(g) illustrate how different divergences learn a student distribution from the given teacher distribution. The a-B3-divergence, compared to others, can more effectively learn soft label information while maintaining focus on the target class.
tions affect the allocation of probability mass in the student distribution during training by tracking the log mass ratio LogR. Notably, LogR is proportional to the gradient of the loss function w.t. the logits. This insight allows to frame the problem as understanding how divergence algorithms influence the reduction of LogR. Through this lens, we identify two key mode-concentratton effects: Hardness- Concentration and Confidence-Conccenttationn Concentration refers to focusing on modes in the loss where there is large error between p and qo, while Confidence- Concentration refers to focusing on modes in the loss where de has high confidence.
a flexible mechanism to smoothly interpolate between the extremes of FKI.D and RKLD by controlling the trade-off between hardness--ooncentration (Fig. 1b) and confidence- concentration (Fig. Ic) via the hyperparameters and B. This mechanism ensures more proper allocation of prob- ability mass (Fig. Ig). Motivated by these insights, we propose ABKD, generic distillation framework based on a- deve-eiieeee Empirical results across a variety of tasks, including instruction-oollwiii and image classification, demonstrate ABKD's generality and effectiveness. For in- stance, by modifying only the loss function, ABKD achieves performance improvements of 0.81 to 3.31 over FKLD and RKLD on five instruction-response datasets when distilling GPT-2 XL (1.5B) into GPT-2 (O.1B).
Concentration. Hardness-
On top of this, we find that the limitations of FKLD and RKLD stem from the extreme ways they utilize these con- centration effects: a) FKLD exhibits weak concentration effects, treating mismatches equally from all classes, which fails in guiding the student to concentrate on the target class and causes incorrect predictions (Fig. Id). b) RKLD hibits strong concentration effects, focusing on both hard classes with large eтrоrs and classes where the student has high confidence. This often leads to a trivial solution, where the well-trained student focuses exclusively on the target class and ignores broader knowledge from p (Fig. le). With the limitations revealed, we continue to seek an answer to the following question:
In summary, the contributions of this work are three-fold:
Theoretically: We analyze the limitations of FKLD and RKLD from novel perspectives of hardness- concentration and confidence-conceetraaiii and show that the a--3-divrrencce offers a flexible approach to balance these effects.
We propose ABKD, flexible distillation framework that unifies FKLD and RKLD and generalizes several other divergences, offering greater versatility and applicability.
Methodologically:
Can we find a generic, theoretically grounded method to balance hardness-conceenratton and confidence- concentration?
Empirically: Extensive experiments on 17 language and vision datasets with 12 teacher-student configura- tions (0.85M-0.44) --0.466 to 7B-3B) validate the theoretical insights. ABKD outperforms or matches state-of-th ohe art methods without extra trainable parameters and allows further gains by rectifying their loss functions.
In pursuit of this, we introduce the a-B-divergence, a gen- eral extension of divergences that unifies FKLD and RKLD, while also extending to previously unexplored divergences like the Hellinger distance and f--diveggence. Our theoret- ical results demonstrate that the a-f--divergence provides
Prior Arts We discuss related work and defer a concen- trated account to App. A.
2

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergence
2. Preliminaries
we uncover the underlying factors that contribute to the lim- itations of FKLD and RKLD by tracking how they allocate probability mass in the student distribution during gradient updates. These insights will guide us in Sec. 4 to identify a more suitable divergence for KD.
KD involves using a fixed teacher model fr to improve the performance of a parameterized student model fs- Given an input æ, the teacher fr and student fs produce probability distributions p and qo, respectively.
The goal of KD can be achieved by letting do mimic p for all samples in dataset D. A direct way to do this is minimizing;
3.1. Tracking Probability Allocation with Log Mass Ratio
(XD (pllgg))
(1)
To find a proper probability matching scheme, KD algo- rithms must keep allocating the probability mass of the student distribution during training. The key in our theory is to keep track of the probability mass change in each gradient update step. To do this, we define monitoring quantity called log mass ratio inspired by Tajwar et al. (2024):
where D is a distribution measure. Optionally, practi- tioners can substitute with the one-hot vector y. where 2 L0....,,...... 0] with at ground-truth label y and O elsewhere. In this case, the loss is lc _ (0))) where ID is typically the FKLD. The final training loss is:
(441(M) g:(v)
LogR," (v) & log
== Ice MLKD.
where A is hyperparameter. Since p provides richer infor- mation (i.e., soft label) than the one-hot vector y. KD out- performs traditional supervised fine-tuning on many down- and image clas-
where g4 (v) is the probability mass for class obtained from algorithm A at step i + 1; (v) is original probability mass for class y at step t.
stream tasks, such as instruction-followe sification. The settings for these tasks in KD are as follows.
To define probability mass qt follow the most popular convention that the class probability is approximated by softmax function such that: q (v) xx exp(ff). where f is the logit for the y-class channel at step t. Interestingly, one can show that (see App. C.1) LogR:" (y) is proportional to the gradient of the logit J as follows:
Instruction-follovingg Let æ and y represent the input and output sequences, respectively. token-level autoregressive model produces an C-dimensional probability distribution for the n-th token over the vocabulary V, conditioned
and yen. where yen ف (y1, V2...... ..n2.))..nn))) 3n-1) denote the generated output sequence up to the (n - th token, The discrepancy between token-level distributions of p and is defined as D(pllg))= A I, Σnz D(p(- æ<<,, (l00( Ven: ))), where T. denotes the sequence length.
LogR (v) = -1 V FTg++N(w))
(5)
where denotes the learning rate and (₂(()) (v) is a normaliz- ing factor independent of y, which vanishes to zero when all the class channel gradients V vanish. Note that un- der a mild assumption, we can show that (see App. C.2) when V W the overall gradient w.r.t. the model weights W, goes to zero, v,, also goes to zero. In this sense, to reach a local minimum, the algorithm automatically reduces the magnitude of v,, f, and also the [LogR (2)1-
Image classification. Let x E H be an image and y CC its one-hot label, with H, W, and C' representing the image dimensions and number of classes. A vision model produces a C-dimensional probability distribution conditioned on æ. The discrepancy between p and qo is defined as (pplgg)).
Based on the discussion above, we next show how reduc- ing (LogR (5) in different divergences affects hardness- concentration and confidence-concenraatio.
3. The Limitations of FKLD and RKLD
Prior arts primarily use FKLD De (pllgo) RKLD
DKr (qellp) to measure distribution discrepancy:
3.2. FKLD and RKLD as Two Extreme Cases
p(k) log p(k)
First, we have the following upper bounds for the log mass ratio for FKLD (Eq. 3) and RKLD (Eq, 4).
DKL = Σ
(3)
go (k))
k
Proposition 3.1. The updates induced by FKLD and RKLD for qs within one gradient descent step are given by:
4o (k)
DxL (golp)) = Σ go (k) log
(4)
pkk)
FKLD:
|LogR (v)l S T 1
(pV)) 1.lv)
qe (v) + +|N (z))
Despite promising success, recent empirical studies show that these two divergences cause suboptimal performance (Wen al., 2023; Ko et al., 2024; Wu et al., 2024). Next,
(a)
(b)
3

---

ABKD:
Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--iivrgeecc
RKLD:
trivial solution: the student concentrates solely the target class and neglects the overall matching effect.
R (1)l S n qc(y) LogR
log p(y) - log (v))
The following theorem makes the above intuition more rig- orous. The results suggest an asymmetric mass allocation for RKLD and an equally important allocation for FKLD. Due to the space limit, the readers are referred to App. G.2 for a formal expression and the proof.
(ai)
(b1)
Σ (c(k))
logp(k) - log qe (x(k))) + | (v),,
(az)
(ba)
where No. (y) and N (y) denote constant normalization
Theorem 3.2 (Informal). Given the student distribution qo and teacher distribution p. FKLD and RKLD differ as follows within one gradient update:
factors independent of y and vanish zero when p 9t-
This proposition is inspired by (Tajwar et al., 2024), but our work distinguishes itself by shifting the focus towards explicitly identifying the factors that drive the decrease in [LogR [v)]- The proof in App. G.1. By minimizing the log mass ratio. there are two types of effects hidden in the results. We first analyze their independent roles.
1. FKLD allocates the mass across all classes equally.
RKLD preferentially increases the mass of underesti- mated (p(y) > qo(y)) classes with higher go(y),
3. RKLD preferentially reduces the mass of overestimated (p(y) < qe(y)) classes with smaller qo
1. The first type, represented by terms (b), (b1), and (b2), has the general form of (s(p()) s(ge (k))), which measures the matching loss between student and teacher distribution. If considered independently, such terms control the effect of ceddnes--ooneenraatinn More precisely, a sharper term with a larger rate-of- change corresponds to an aggressive student who aims to focus on the hardest classes reach a good match- ing performance w.E.t. the teacher (Fig. 1b).
As shown in Fig. I(d), the equally weighted matching scheme in FKLD drives students to sub-optimal modes, which induces wrong predictions. For RKLD, the theorem states that it only favors small mass classes when the teacher score is over-estimated while only favors large mass classes under the opposite scenario. As a total effect, the small mass tends to get smaller; the large mass tends to get larger. As an extreme result shown in Fig. 1(e), RKLD eventually forces the student focus on one class. This makes the teacher's supervision degenerate to a ont-hot la- bel, which loses the distributional information hidden inside the teacher's prediction. This leads to the following conclu- sion:
2. The second type, denoted by terms (a), (a1). (a2). can be expressed the student's confidence weight- ing function: qe (1) (3 > 0). If considered indepen- dently, such terms control the effect of confidence- concentration. In other words, sharper weighting function corresponds to a confident student who only cares about the matching performance in classes that the student believes to be the ground truth (Fig. Ic).
A proper divergence should achieve a moderate trade- off between hardness-cooncenttation and confidence- concentration.
What is the joint effect of the two? Prop. 3.1 provides two extreme answers, FKLD and RKLD. FKLD in the re- sults picks a very weak hardness-concentrann s(x:) and very weak concentration effect with effect with 3 0. As result, FKLD forces the student to treat all matching penalties equally for all classes since there is no weighting function ((a) 1). This fails to concen- trate on the target classes. By contrast, RKLD picks a very strong hardness-concentratie
3.3. Weighted Sum of FKLD and RKLD
In pursuit of this, naive solution is to take a weighted sum of FKLD and RKLD, which we call the weighted sum divergence (WSD):
Dwws(plla) e aD (plla) + 3IK.. (q [p),
(6)
effect with s(x) = log(x) (recall that 0 x < 1. log is much sharper than linear function) and a very strong confidence--coneetraaiie with 3 = 1. Recall that a well-trained student distribu- tion 4: primarily concentrates probability on the target class and assigns smaller probabilities to others. In this case, an overly strong conceettctcoonnttttiiioooooo effect suppresses the effect on non non-target classes while emphasizing this effect on the target class. This results in
where a and ,3 are hyperparameters. A more principled approach is adapt the weighting coefficients dynamically during training based the discrepancy (e.g., entropy) between p and g, as done in previous works (Amara et al., 2022; Wu et al., 2024).
effect
Unfortunately, such a composite metric overemphasizes modes with small probabilities in p and To see this, when either q(k) Ri 0, (k)) 0 or p(k) z O, (ik) > O, we have
4

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergence
(wsD(Plly) oc. Hence, the algorithm must focus on extreme cases to minimize the objective function, leading to improper probability allocation. Moreover, similar to the analysis in Ko et al. (2024), one can easily show that the gradient norm in this case also grows excessively, lead- ing to significant and potentially noisy parameter updates. Such behaviors can destabilize the optimization process and hinder convergence.
Table 1. Some divergence functions and their corresponding choices of and B. The a-3-divergence can he extended by continuity (by applying l'Hópital formula) to cover all the values of a, 3 E R. as shown App. B.
Distribution Measure
Reference
Runge
Kullback-Leibler (KL) divergence Reverse KL divergence
Kullback & Leibler (1951)
1. 8
Kullback & Leibler (1951)
a = 0,
c-divergeace
Chernoff (1952)
a+ -+
d-divergence
Basu et (1998)
Hellinger distance
Helllngee (1909)
a a= 8 0.5
Another attractive approach is to use the Jensen- -hannon divergence (Binici et al., 2022; Agarwal et al., 2024) (isD(pllq) & 1DLL (p) [m) + DRL (q| (|lm) where m = [P + q). However, a major drawback of JSD is that it suf- fers from gradient vanishing (Arjovsky et al., 2017) when the distributions p and qo are far apart (a common scenario in early training stages), which hinders model convergence.
Squared enclidean distance
Heath (1956)
0
of a family distribution divergences, which includes FKLD, RKLD, and other typical divergences as special cases. For example, when (a = 1, B = 0). one obtains FKLD; when (ar = 0, = 1), one obtains RKLD. Please see Tab. 1 for other special cases. In this way, ABKD nature provides a generic framework for divergence-basel distillation algo- rithms.
Above all, balancing hardness and confidence concentra- tion is non-trivial if one only resorts to FKLDD and RKLD.. In the next section, we will introduce a generic notion of divergence to address this issue.
4.2. Trading off Harddees-CConcentrude and Confidence-Concentraaie via a-33divergence
4. ABKD: The Proposed Method
ABKD offers unified space to trade off the hardness- concentration and confidence-concennee concentration effects.
4.1. ABKD
One way to pursue harmonic utilization of hardness- and confideece--conenntatiin is to find subtle point between FKLD and RKLD. The following a- ----iivereeece exactly serves this purpose (Cichocki et al., 2011).
To explain this, we go back to the log mass ratio, the follow- ing proposition explains how the hyperparameters a and influence the reduction of [LogR (a.0) (v)l-
Proposition 4.2. The updates induced by a- ----diiergeece for qe within one gradient descent step are given by:
Definition 4.1 (o---divergvencee. Consider a and 3 E RR satisfying a, B, a + 3 + O. the a-S-divergence of two distributions given by:
p(y) ge(
|oggR (w))l S n
() D 0
(p 11
g)))
Σ a.
1
p(k)''(())
a a pp())att
(a)
(b)
a + (k))ttt
3
+nde(y) Σ (k)
p(k) n - 90 (k))
+ +NN(a.9)(0)
a
(01)
(b1)
where p- = [o(k) and q = ( are two discrete
distributions over C' classes.
where No. (y) denotes constant normalization factiiiaaaiooffoooooooo inde- pendent of y and vanishes to zero when p = q:-
As will soon be seen in Sec.4.2, both hardness- concentratonn and confidence-concenttatiinn effects in a-f3-divergence could be regarded as an interpolation between the corre- sponding effect of FKLD and RKLD. Such ability allows the a-33-divergence to ensure more proper allocation of probability mass.
The proof is in App. G.5. In (a) and (a1), a-3- divergence employs a power form (k)) for confidence- concentration effect. It is easy see when B - 1, it degenerates to the effect of RKLD, and when 3 - 0 to the effect of FKLD. A larger 3 provides a stronger ef- fect of conddeece-concentatiionn focusing the matching performance on its most confident classes (Fig. Ic). Mean- while, terms (b) and (b1) uses p(v)"= for hardness- concentration effect. is easy to see when a - 1, it degen- crates to the effect of FKLD, and when a + 0 to the effect of RKLD. smaller a amplifies the hardness-concentee
Inspired by this, we propose ABKD, which is formally defined as minimizing the following objective:
(= + (a.8) (p go):
(7)
Beyond this issue, a- B-vvviieeeee also a generic notion
5

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergence
Table ROUGE-L scores (1) on five task-agossi instruction followiigg datasets. Note that this an unfair comparison because we only train on the fixed dataset while other KD methods employ augmentation. The fairer results using our method with different augmentation strategies can be found in Fig. 3(b) and Tah. 8. All results are based on our re-implementation. We report the average standard deviation of ROUGE-L scores across five random seeds. Better results are shown in bold, and darker colors indicate superior performance.
Method
Dolly Eval
Self-Instruct
Vicuna Eval
Super-Natural
Unnatural
GPT-2 XL (Teacher)
26.94 (0.23)
13.31 (0.63)
16.23 (0.62)
24.28 (0.43)
29.05 (0.14)
GPT-2 XL (L.5B) GPT-2 (0.1B)
SFT
23.14 (0.23)
10.22 (0.44)
15.15 (0.31)
17.41 (0.18)
19.76 (0.09)
KD (Himton, 2015)
23.80 (0.37)
10.01 (0.75)
15.25 (0.65)
17.69 (0.26)
18.99 (0.05)
SeqKD (Kim & Rush, 2016)
24.28 (0.22)
11.24 (0.30)
14.94 (0.58)
20.66 (0.28)
23.59 (0.13)
MiniLLM (Gu al., 2024a)
24.62 (0.33)
12.49 (0.56)
17.30 (0.41)
23.76 (0.38)
24.30 (0.14)
GKD (Agarwal et al., 2024)
24.49 (0.16)
11.41 (0.14)
16.01 (0.37)
18.25 (0.24)
21.41 (0.11)
DISTILLM (Ko et al., 2024)
25.32 (0.14)
11.65 (0.28)
16.76 (0.66)
23.52 (0.47)
25.79 (0.08)
Ours (ABKD)
25.65 (0.24)
13.47 (0.42)
16.06 (0.25)
26.47 (0.31)
29.32 (0.08)
GPT-2 XI. (1.58) - GPT-2 Medium (0.3B)
SFT
25.30 (0.31)
12.56 (0.62)
16.36 (0.22)
23.32 (0.13)
23.42 (0.07)
KD (Hinton, 2015)
24.71 (0.17)
10.33 (0.54)
16.23 (0.50)
23.74 (0.32)
23.97 (0.12)
SeqKD (Kim & Rush, 2016)
25.93 (0.44)
12.98 (0.24)
16.68 (0.30)
21.95 (0.19)
25.23 (0.08)
MiniLLM (Gu et al., 2024a)
25.34 (0.25)
13.36 (0.62)
17.25 (0.46)
25.68 (0.41)
26.63 (0.12)
GKD (Agarwal et al., 2024)
24.75 (0.27)
12.76 (0.85)
16.54 (0.39)
24.94 (0.14)
26.42 (0.15)
DISTILLM (Ko et al., 2024)
26.21 (0.29)
13.53 (0.13)
16.96 (0.66)
25.78 (0.19)
28.51 (0.26)
Ours (ABKD)
26.08 (0.36)
13.86 (0.40)
16.63 (0.26)
27.25 (0.38)
29.69 (0.21)
GPT-2 XL (I.5B) + GPT-2 Large (0.8B)
SET
25.42 (0.32)
12.91 (0.46)
16.31 (0.51)
23.76 (0.28)
25.72 (0.07)
KD (Hinton, 2015)
26.02 (0.43)
12.34 (0.52)
16.26 (0.44)
25.11 (0.37)
26.44 (0.12)
SeqKD (Kim & Rush, 2016)
26.29 (0.47)
13.53 (0.34)
16.39 (0.36)
25.81 (0.40)
27.51 (0.10)
MiniLLM (Gu et al., 2024a)
26.12 (0.25)
13.79 (0.31)
17.35 (0.51)
26.12 (0.37)
28.53 (0.17)
GKD (Agarwal et al., 2024)
26.06 (0.34)
13.21 (0.45)
16.64 (0.45)
26.13 (0.41)
27.13 (0.21)
DISTILLM (Ko et al., 2024)
26,56 (0.36)
13.97 (0.36)
16.61 (0.45)
26.73 (0.36)
29.24 (0.23)
Ours (ABKD)
26.51 (0.22)
14.38 (0.43)
16.63 (0.42)
28.05 (0.21)
29.92 (0.14)
effect and tends to be more aggressive in achieving better matching by penalizing errors on hard classes (Fig. 1b).
in Fig. ((a) and (f), For further analysis, please see App. F.
In this sense, by tuning a and B, we can flexibly balance the influence of the two effects and avoid extreme cases (Fig. 1g). For a finer-gu merrgrraiee theoretical analysis and hyper- parameter tuning guidelines, please see App. D, Thm. D.1.
5. Experiments
In the following, we investigate to what extent our theoret- ical results translate into practice natural language and vision tasks. Due to space limitations, please see App. I for more details on datasets, competitors, and implementation.
Comparing with the Weighted Sum. As discussed earlier in Sec. 3.3, WSD often focuses excessively on the extreme values of p/o g, leading to unstable optimization. Fortunately, one can show that the a----divergence can finely adjust the focus different likelihood ratios p/q, thus enjoying more stable gradient. For further analysis, see App. E.
5.1. Natural Language Processing Tasks
Datasets. We evaluate our methods on five task-agnostic instruction-following benchmarks. Evaluation metric is based on ROUGE-L (Lin, 2004). Details about the datasets and evaluation metric can be found in App. I.I.1.
Comparing with the a-divergence. One might also recall the a-divergence to achieve the trade-off, which is defined as D& (pllg) 2 I] [E, ((k)) - 1]. It includes DeL (pllg) as - 1, and DEL (ye (p) as - 0. Note that when / = - a, it becomes a special case of our framework. According to Prop. 4.2, to decrease a, one has to increase 3 to ensure that they add up to 1. Such unnecessary restriction hinders its ability to achieve better performance, as shown
Competitors. We consider the following state-of of-the-art (SOTA) baselines: 1) supervised fine-tuning (SFT) with only student model on fixed datasets; 2) KD with FKLD on fixed datasets; 3) SeqKD with SFT to macherrgeneeaal 4) MiniLLM with RKLD using a policy gradient approach on student-generated outputs (SGOs): 5) GKD with JSD on
output;
6

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via 1e-3--Diveggece
Table Evaluation of effect different loss functions. WSD: weighted sum of FKLD and
28.0
RKLD. HD: Hellinger distance. SED: Squared euclidean distance.
Val ROUGE-L
26.0
Loss Function
Dolly Eval
Self-Instruct
Vicuna Eval
Super-Natural
Unnatural
24.0
FKLD
23.80 (0.37)
10.01 (0.75)
15.25 (0.65)
17.69 (0.26)
18,99 (0.05)
RKLD
24.77 (0.37)
12.02 (0.48)
15.06 (0.28)
23.27 (0.29)
26.01 (0.11)
22.0
WSD
23.33 (0.52)
10.52 (0.47)
14.83 (0.61)
19.67 (0.13)
21.21 (0.21)
HLLD
560
HD
25.15 (0.36)
12.39 (0.77)
15.43 (0.20)
24.14 (0.23)
26.83 (0.15)
SED
21.04 (0.51)
10.00 (0.56)
13.73 (0.17)
19.34 (0.19)
22.62 (0.19)
10 epoch
20
a-divergence
25.15 (0.41)
12.92 (0.22)
15.60 (0.27)
24.83 (0.21)
27.81 (0.10)
3-divergence
24.12 (0.38)
11.18 (0.27)
14.95 (0.33)
20.98 (0.23)
23.15 (0.14)
Figure 2. Performance across different
a-B-divergence
25.65 (0.24)
13.47 (0.42)
16.06 (0.25)
26.47 (0.31)
29.32 (0.08)
loss functions the validation set.
29.2
Samples/sec
2.5
Val ROUGE-L
29.0
87.5
2.0
85.0
Accuracy (%)
1.5
82.5
A
A
1.0
80.0
A
28.4
A New
77.5
28.2
40
Sepro
6x0
DETALLL
Duce
Fladd
Omcpolicy
Mixed
Adopive
75.0
72.5
(a) Training Speed
(b) Effects of SGOs
70.0
Figure 3, Comparison of training speeds and the effects of using SGOs. Please see Sec. I.1.2 for details of different SGOs strategies.
Teacher Ours
LSD
DKD
KD
PromptsRt
MaPLe
CoCoOp
CLIP
Figure 4. Comparison with SOTA methods on base-to-new setting, HM denotes the harmonic mean of base and new accuracy. Results are averaged across datasets, with per-dataset details in Tab. 16. Results of haseline CL.IP are evaluated on the pre-trained model. Teacher: ViT-I/14 CLIP; Student: ViT-B/16 CLIP.
mixture of SGOs and fixed datasets; 6) DISTILLM with S(R)KL mixture of SGOs and fixed datasets. Please refer to App. 1.1.2 for details about competitors and SGOs.
Results. From Tab. 2, we have the following observations: 1) Distillation methods often outperform SFT, showcasing their potential. However, they can sometimes yield worse re- sults (e.g. KD on Unnatural when distilling GPT-2 XL into GPT-2), highlighting the importance of selecting a proper distillation objective. 2) By simply modifying the distil- lation objective, our framework outperforms vanilla KD and SFT across various datasets when distilling GPT-2 XL. (I.5B) to smaller-scale families of GPT-2 (0.1B~0.8B): Prior arts (Ko et al., 2024; Agarwal et al., 2024) show that training with SGOs can lead to significant improvements. However, even when compared to KGOs-based methods (e.g., GKD, DISTILLM) under this inherently unfair setting, our approach consistently achieves superior or comparable results, especially on Super-Naturll and Unnatural datasets.
proaches. As shown in Fig. 3(b), our framework consis- tently delivers high performance across different settings, highlighting its adaptability and effectiveness. Further ex- perimental results and analyses are provided in App. 1.1.1.
Effects of Loss Functions. Tab.3 compares the perfor- mance between various loss functions. The results show that a-8-divergence consistently outperforms the others, while using only a- or 3-divergence degrades performance due to limited expressivity. achieves improvements of 0.81 1 to 3.31 over FKLD and RKLD DD across five datasets, whereas WSD, which combines weighted FKLD and RKLD, fails to deliver comparable re- sults. Furthermore, Fig.2 demonstrates the superior perfor- mance of a- a-B-ddvergence during the entire training phase.
particular, a- a-8--divergence
Efficiency Comparison. Fig. 3(a) shows that our frame- work matches the training speed vanilla KD, as only modifies the distillation objective without introducing ad- ditional cost. This addresses concerns regarding the scala- bility of our method. In contrast, other distillation methods require 1.6 to 7 times longer training time due to the continuous need to sample student's outputs during training.
In summary, these empirical results align with the theoretical insights in Sec. and show that even modest adjustments to the loss function can yield significant improvements.
5.2. Vision Tasks
Datasets. We conduct experiments on 12 popular im- age recognition datasets. Dataset details are referred to App. I.2.1. The evaluation metric used is accuracy. Apart
Effects of SGOs. We examine the robustness of our frame- work by evaluating performance with various SGOs ap-
7

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via --e----Dieeeenee
77.0
76.0
76.5
72.0
71.5
yareeacyyy 0%)
75.0
73.5
Accuracy
71.0
70.5
Accuranng
23.3
73.0
69.0
TMM
ABDIO
L20
AU1D
Marcher
Neacher
ABLSD
ABDIO
Dko
4310
150
Teacher
AB7YN
Dka
T7W
ABKO
necher
ABLSD
407TH
ABDKO
ADKO
650
DK0
TTM
(a) WRN-40-2 WRN-16-2
WRN-40-2
WRN-40-1
resnet56
resnet20
resnet 10
resnet 20
76.0
74.0
23.3
76
qut
Accuracy
23
2M
Accuracy
75.5
75.0
74.5
74.0
72.0
72
73.0
71.5
72.5
72.0
ABTTI
Tencher
ADDNO
480
Reacher
ABLID
ABOVO
150
DK0
T7M
ADKD
80
arf
ABDED
ARTrI
DkD
Rarther
TM
ABKO
80
ABLSO
ABDYD
ABVo
DKD
TM
(e) resnetl 10
resnet32
32x4
resnet8x4
vgg
vgg8
(h) resnet
resnet44
Figure 5. Accurayy CIFAR- 100 for student models trained with different distillation methods. ABDKD, ABTTM, ABLSD, and ABKD are our implementations by rectifying the backbone's 's loss function. For details on backbones, please refer to Sec. I.2.2.
3.78
40
Train Accuracy (%)
60
ROUGE-L
Entropy
Self-BLEU
0.p
3.74
30
40
a.pha=0.2
alpha-0.2
0.5
alp
0.7
0.7
1.0
alpha=1.0
3.70
20
10
50
100
200
250
10
0.0
0.7
1.0
0.5
0.7
1.0
epoch
epoch
beta
beta
(a) alpha on CIFAR- 100
(b) alpha on Dolly
(c) beta on CIFAR-100
(d) beta on Dolly
Figure 6. Sensitivity analysis of hyperparameters and 8. (a)-(b) For low-dimensional output distribution in CIFAR-100, smaller leads to excessive penalization for error with limited gains. However, for higher-dimensional distrihution in Dolly (egg. 50,527 for GPT-2), well-tuned smaller critical. (c)-(d) A larger sharpens output distributions by emphasizing classes with high student confidence.
from the standard training-evaluatio paradigm, we further consider novel base-to-new setting (Zhou al., 2022; Hua thoroughly analyze the student model's
our framework selects more suitable distillation objectives for specific teacher-studens pairs, showing competitive or superior results, particularly in base-to-new setting.
al., 2025) generalization across classes. In this setup, training is per- formed on base classes, and accuracy is evaluated on both base and new classes. Please see App. I.2.3 for more details.
Apply to Other Distillation Techniques. Fig. 5 also shows that our framework can act as simple plug-and-pla tool to rectify the loss functions used by existing methods, yielding further improvements (e.g.. ABDKD vs DKD).
Competitors. We consider the following SOTA distillation methods: 1) KD, 2) DKD, 3) LSD, and 4) TTM. For the base-to-new setting, we also compare with SOTA SFT meth- ods: 5) CoCoOp. 6) MaPLe, and 7) PromptSRC. Please refer to App. I.2.2 for method details.
5.3. Sensitivity Analysis
We next analyze the effects of ceadness-conccenrattion and fifidence-connenrationn which helps validate the theoreti- cal insights shown in Prop. 4.2 and Thm. D.I.
Results. Fig. 4 and Fig. 5 show results from 9 teacher- student architectures on 12 datasets. Based on these, we conclude: 1) Without modifying the distillation objective, methods that more effectively utilize teacher distribution knowledge (e.g., DKD, TTM, and LSD) can outperform vanilla KD; 2) However, their scores fall short in some cases, such as LSD in base-to-new setting; 3) Orthogonal to them,
Figs. 6(a) and (b) show performance during training for different a. In CIFAR- 100, with its relatively dinddimnsional output distribution, smaller (stronger hardness-conneetraaion aggressively penalizes errors but offers limited gains. However, in Dolly,
Effect of a on
8

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergenc
with higher-dimeensonaal output (e.g., GPT-2's vocabulary size of 50,257), well-tuned smaller is crucial to avoid local optima, especially in early training stages.
potential has been demonstrated in some recent frontier works, such DeepSeek-R1 (Guo et al., 2025) and Qwen- 3 (Yang et al., 2025). Although these methods primarily adopt SeqKD- distillation techniques (Kim & Rush, 2016) and therefore do not involve distribution matching through KL divergence, the results in Tab. 2 demonstrate that logit- methods possess even greater potential. We believe that the applications of KD techniques will continue to expand.
Figs. 6(c) and (d) show how 3 affects Shannon entropy of the output dis- tribution and Self-BLEU score (Zhu et al., 2018) of output sequences (100 indicates deterministic outputs and 0 denotes maximum diversity). The smaller 3 (weaker confidence- concentration) places more emphasis on classes with low student confidence, encouraging the student to focus more on learning the soft label information from the teacher distri- bution. This leads to a smoother output distribution (higher entropy) and more diverse generated sequences (lower Self- BLEU). Thus, selecting an appropriate 3 ensures a balance between focusing on the target class and learning more soft label information.
Effect of B on confidence--ooneettaatien
In this work, we provide a theoretical analysis of the fun- damentally different behaviors-mode-cooeeerr mode-covering and mode- seeking-
by forward and reverse KL divergences in KD, grounded in two novel mode-concentratioo effects. These new theoretical insights further shed light on the development of more principled and effective distillation objectives. Interestingly,
theoretical framework also aligns with findings from related studies, such as the phe- nomenon of likelihood displacement (Razin et al., 2024; Ren & Sutherland, 2024) observed in Direct Preference Op- timization (DPO). This phenomenon stems from the equiv- alence of DPO w.r.f. the reverse optimization of KI. diver- gence (Rafailov et al., 2023; Tajwar et al., 2024).
6. Conclusion
In this paper, we argue that the key to KD lies in trading off two mode-concentration effects: and monfidence-conennraati
The widely used FKLD and RKLD fail achieve this balance, instead representing two extreme cases that lead to improper probability alloca- tion. To address this issue, we introduce ABKD, a generic distillation framework based on a-8--iiergence. ABKD generalizes FKI.D and RKI.D to a broader family of diver- gences, offering greater flexibility. Our theoretical results show that ABKD can flexibly interpolate between the above two extremes, enabling an effective trade-off. Extensive experiments further demonstrate its effectiveness.
References
Agarwal, R., Vieillard, N., Zhou, Y., Stanczyk, P., Garea, S. R., Geist, M., and Bachem, O. On-policy distillation of language models: Learning from self-generater mistakes. In The Twelfth International Conference on Learning Representations, 2024.
Amara, I., Sepahvand, N., Meyer, B. H., Gross, W. J., and Clark, J. J. Bd-kd: balancing the divergences for online knowledge distillation. arXiv preprint arXiv:2212.12965, 2022.
Acknowledgements
This work was supported in part by the Fundamental Research Funds for the Central Universities, in part by the National Key R&D Program of China under Grant 2018AAA0102000, 2024QY210004 and 2022YFC3302300, in part by National Natural Science Foundation of China: 62236008, 62441232, U21B2038 U23B2051, 62122075, 62206264 and 92370102, in part by Youth Innovation Pro- motion Association CAS, in part by the Strategic Priority Research Program of the Chinese Academy of Sciences, Grant No.XDD0680201, in part by the China National Postdoctoral Program for Innovative Talents under Grant BX20240384. The authors thank anonymous reviewers for their insightful comments and suggestions.
Arjovsky, M., Chintala, S., and Bottou, L.. Wasserstein gen- erative adversarial networks. In International conference machine learning, pp. 214-223. PMLR, 2017.
Basu, A., Harris, I. R., Hjort, N. L., and Jones, M. Robust and efficient estimation by minimising density power divergence. Biometrika, 85(3):44--55) 1998.
Binici, K., Pham, N. T., Mitra, T., and Leman, K. Prevent- ing catastrophic forgetting and distribution mismatch in knowledge distillation via synthetic data. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 663-671, 2022.
Impact Statement
Bossard, L., Guillaumin, M., and Van Gool, L. Food-101- mining discriminative components with random forests. In Computer vision-ECCV 2014: 13th European con- ference. zurich, Switzerland, September 6-12, 2014, pro- ceedings, part VI 13, pp. 446 446-461. Springer, 2014.
The work presented in this paper aims to advance the field of knowledge distillation (KD), promising direction to enable knowledge transfer between different models. This

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divrggene
Chen, D., Mei, J.-P., Zhang, H., Wang, C., Feng, Y., and Chen, C. Knowledge distillation with the reused teacher classifier. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11933- 11942, 2022.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Heath, L. The thirteen books of Euclid's Elements. Dover Publications, Inc, 1956.
Chernoff, H. A measure of asymptotic efficiency for tests of a hypothesis based on the sum of observations. The Annals of Mathematical Statistics, pp. 493-507, 1952.
Helber, P., Bischke, B., Dengel, A., and Borth, D. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Se- lected Topics in Applied Earth Observations and Remote Sensing, 12 2019.
Cho, J. H. and Hariharan, B. On the efficacy of knowledge distillation. In Proceedings of the IEEE/CVF interna- tional conference computer vision, pp. 4794-4802, 2019.
Hellinger, E. Neue begründung der theorie quadratischer formen von unendlichvielen veränderlichen. Journal für die reine und angewandle Mathematik, 1909(136):100-
Cichocki, A., Cruces, S., and Amari, S.-i. Generalized alpha-beta divergences and their application robust nonnegative matrix factorization. Entropy, 13(1):134- 170, 2011. ISSN 1099-4300. doi: 30.3394 htthttp::/ //tp:://wwwww// htt/:/wwwwpp....m//////////// 1/134.
271, 1909.
Hinton, G. Distilling the knowledge in neural network. arXiv preprint wrxiv:1503.0253/ 2015.
Honovich, O., Scialom, T., Levy, O., and Schick, T. Unnat- ural instructions: Tuning language models with (almost) no human labor. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61 Annual Meet- ing of the Association for Computational Linguistics (Vol- ume Long Papers), pp. 14409- 14428, Toronto, Canada, July 2023. Association for Computational Linguistics. URL https:
Conover, M., Hayes, M., Mathur, A., Xie, J., Wan, J., Shah, S., Ghodsi, A., Wendell, P., Zaharia, M., and Xin, R. Free dolly: Introducing the world's first truly open instruction- tuned 1lm. Company Blog of Databricks, 2023.
Deng, J., Dong, W., Socher, Li, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Iece, 2009.
doi: 10. 10.116653/11/2233ac--oomoo / //aclant wwtp//23../..... acl long. 806.
Hua, C., Xu, Q.. Yang. Z., Wang, Z., Bao, S., and Huang, Q. Openworldauc: Towards unified evaluation and opti- mization for open- open-world prompt tuning. arXiv preprint arXiv:2505.06 05 180, 2025.
Dosovitskiy, A. An image is worth 16x16 words: Trans- formers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Fei-Fei, L., Fergus, R., and Perona, P. Learning generative visual models from few training examples: An incremen- tal bayesian approach tested on 101 object categories. In 2004 conference on computer vision and pattern recogni- tion workshop. pp. 178-178. IEEE, 2004.
Huang, T., You, S., Wang, F., Qian, C., and Xu, C. Knowl- edge distillation from a stronger teacher. Advances in Neu- ral Information Processing Systems, 35:33716-33721 2022.
Huang, 7. and Wang, N. Like what you like: Knowl- edge distill via neuron selectivity transfer. arXiv preprint arXiv:
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open- webtext corpus, 2019.
1tf.:::777.012/111 2017.
Gu, Y., Dong. L., Wei, F., and Huang, M. Minillm: Knowl- edge distillation of large language models. In The Twelfth International Conference on Learning Representations, 2024a.
Jiao, X., Yin, Y., Shang, L., Jiang, X., Chen, X., Li, L., Wang, F., and Liu, Q. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv: /909. 10351, 2019.
Gu, Y.. Zhou, H., Meng, F., Zhou, J., and Huang, M. Mini- plm: Knowledge distillation for pre-training language models. arXiv preprint arXiv:2410.17215. 2024b.
Jin, Y., Wang, J., and Lin, D. Multi-level logit distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 24276-24285, 2023.
Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., et al. Deepseek-rl: In- centivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501. 12948, 2025.
Khattak, M. U., Rasheed, H., Maaz, M., Khan, S., and
Khan, F. S. Maple: Multi-modal prompt learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. pp. 19113-19122, 2023a.
10

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3-Divergence
Khattak, M. U., Wasim, S. T.,, Naseer, M., Khan, S., Yang. M.-H., and Khan, F. S. Self-regulating prompts: Foun- dational model adaptation without forgetting, In Pro- ceedings of the IEEE/CVF International Conference on Computer Vision, pp. 15190-15200, 2023b.
Lin, C.-Y. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
Lv, J., Yang, H., and Li, P. Wasserstein distance rivals kullback-leibler divergence for knowledge distillation. arXiv preprint arXiv:2412.08139 2024.
Kim, G., Jang, D., and Yang, E. PromptKI: Dis- tilling student-friendly knowledge for generative lan- guage models via prompt tuning. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Find- ings of the Association for Computational Linguis- tics: EMNLP 2024, pp. 6266-6282, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 1.11866/1112222.nnniiii 364. URL https:ttp////.cctc /t///claanttt www.anhollggy.o.ooooooooooololaa org/2044. Findings-emni 3.34//
Maji, S., Rahtu, E., Kannala, J., Blaschko, M., and Vedaldi, A. Fine- aene-grained visual classification of aircraft. arXiv preprint arXiv:130..3151, 2013.
Mirzadeh. S. I,, Farajtabar, M., Li, A., Levine, N., Mat- sukawa, A., and Ghasemzadeh, H. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5191-5198, 2020.
Kim, Y. and Rush, A. M. Sequence-level knowledge dis- tillation. CoRR, 00946606777777 2016. URL http: //arxiv. wwwp/vvvoorg/io.o abs 166.. 07947.
Nilsback, M.-E. and Zisserman, A. Automated flower clas- sification over large number of classes. In 2008 Sixth Indian conference on computer vision, graphics & image processing, pp. 722-729. IEEE, 2008.
Ko, J., Kim, S., Chen, T., and Yun, S.-Y. Distillm: Towards streamlined distillation for large language models. arXiv preprint arXiv:2402.00 03898, 2024.
Ouyang, L., Wu, J., Jiang. X., Almeida, D., Wainwright, C., Míshkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., al. Training language models follow instructions with human feedback. Advances in neural information processing systems, 35:27730-2775 2022.
Krause, J., Stark, M., Deng. J., and Fei-Fei, L. 3d object rep- resentations for fine- fine-graine grained categorization. In Proceed- ings of the IEEE international conference on computer vision workshops. pp. 554-561, 2013.
Parkhi, O. M., Vedaldi, A., Zisserman, A., and Jawahar, C. Cats and dogs. In 2012 IEEE conference on computer vision and pattern recognition, pp. 3498-3505. IEEE, 2012.
Kullback. S. and Leibler, R. A. On information and suf- ficiency. The annals of mathematical statistics, 22(1): 79-86, 1951.
Lee, H., Park, Y., Seo, H., and Kang, M. Self-knowledge distillation via dropout. Computer Vision and Image Understanding, 233:103720, 2023.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763, PMLR, 2021.
Li, L., Bao, Y., Dong, P., Yang, C., Li, A., Luo, W., Liu, Q. Xue, W., and Guo, Y. Detkds: Knowledge distillation search for object detectors. In Forty-first International Conference on Machine Learning, 2024.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Er- mon, S., and Finn, C. Direct preference optimization: Your language model is secretly reward model. Ad- vances in Neurul Information Processing Systems, 36: 33728-53741, 2023.
Li, X.-C,, Fan, W.S., Song, S., Li, Y., Yunfeng, S., Zhan, D.-C., et al. Asymmetric temperature scaling makes larger networks teach well again. Advances in neural information processing systems, 35:3830-3842, 2022.
Razin, N., Malladi, S., Bhaskar, A., Chen, D., Arora, S., and Hanin, B. Unintentional unalignment: Likelihood displacement in direct preference optimization. arXiv preprint 1rXiv:2410.08847, 2024.
Liang, C., Jiang. H., Li. Z., Tang, X., Yin, B., and Zhao, T. Homodistil: Homotopic agn-aggnstic distillation of pre- trained transformers. arXiv preprint arXiv:2302.00663. 2023a.
Ren, Y. and Sutherland, D. J. Learning dynamics of 11m finetuning. arXiv preprint arXi 2024.
Liang, C., Zuo, S., Zhang, Q., He, P., Chen, W., and Zhao, T. Less more: Task- layer-wiss distillation for lan- guage model compression. In International Conference on Machine Learning, pp. 2065-70067. PMIR, 2023b.
Romero, A., Ballas, N., Kahou, S. E., Chassang, A., Gatta, C., and Bengio, Y. Fitnets: Hints for thin deep nets, arXiv preprint 4 141.6555 2014.

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergence
Simonyan, K. Very deep convolutional networks for large- scale image recognition. arXiv preprint arXiv: 2014.
Wang, Y., Mishra, S., Alipoormolabashi, P., Kordi, Y., Mirzaei, A., Naik, A., Ashok, A., Dhanasekaran, A. S., Arunkumar, A., Stap, D., Pathak, E., Karamanolakis, G., Lai, H., Purohit, I., Mondal, I., Anderson, J., Kuz- nia, K., Doshi, K., Pal, K. K., Patel, M., Moradshahi, M., Parmar, M., Purohit, M., Varshney, N., Kaza, P. R., Verma, P., Puri, R. S., Karia, R., Doshi, S., Sampat, S. K., Mishra, S., Reddy A, S., Patro, S., Dixit, T., and Shen, X. Super-Naturallnnstuciinn Generalization via declar- ative instructions on 1600+ NLP tasks. In Goldberg, Y., Kozareva, 7.,, and Zhang. Y. (eds.). Proceedings of the 2022 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 5085-5100 Abu Dhabi, United Arab Emirates, December 2022a. Association for Compu- tational Linguisties. doi: 10.11665//1/2022..co main. 340. URL https: / / http///aaaaathllooooo .rrg:/org///// 2022. emnlp-maine 340.
Soomro, K. Ucf101: A dataset of 101 human actions classes from videos in the wild. arXiv preprint arXiv:/212.0400) 2012.
Sun, S., Cheng, Y., Gan, Z.,, and Liu, J. Patient knowledge distillation for bert model compression. arXiv preprint arXiv: MeXivv/908.000 09355, 2019.
Sun, S., Ren, W., Li, J., Wang, R., and Cao, X. Logit standardization in knowledge distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 15731-15740, 2024.
Sun, Z., Yu, H., Song, X., Liu, R., Yang, Y., and Zhou, D. Mobilebert: a compact ausk-agnostic bert for resource- limited devices. arXiv preprint 20X1v:2004.02984, 2020.
Wang. Y., Kordi, Y., Mishra, S., Liu, A., Smith, N. A., Khashabi, D., and Hajishirzi, H. Self-instruct: Align- ing language models with self-generated instructions. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Pro- ceedings of the ó1 Anmual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), pp. 13484-13508, Toronto, Canada, July 2023. Associ- ation for Computational Linguistics. doi: 10. 18.18v53/v/ 2023
Tajwar, F., Singh, A., Sharma, A., Rafailov, R., Schneider, J. Xie, T., Ermon, S., Finn, C., and Kumar, A. Preference fine-tuning of lims should leverage suboptimal, on-policy data. arXiv preprint arXiv:2404. 14367, 2024.
Tian, Y., Krishnan, D., and Isola, P. Contrastive representa- tion distillation. arXiv preprint arXiv: 10699, 2019.
URL https::/p//oo. lnttholoyy. /rg/2023-02 aew- 754.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F,, et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint 2rXiv:2300. 2023a.
Wang. Z., Xu, Q., Yang, 7., He, Y., Cao, X., and Huang, Q. Openauc: Towards auc-oriented open-set recognition. Advances in Neural Information Processing Systems, 35: 25033-25045 2022b.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A.. Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models, 2023b. URL https: //arii/aaviivaai org/ abs/ 13971.
Wen, Y., Li, Z., Du, W., and Mou, L. f-divergence mini- mization for sequence-level knowledge distillation. arXiv preprint arXiv:2307.15199 2023.
Wu, T, Tao, C., Wang. J., Yang, R., Zhao, Z., and Wong, N. Rethinking kullback-leibler divergence in knowledge distillation for large language models. arXiv preprint MoXi::244.002677, 2024.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wang, D., Gong, C., Li, M., Liu, Q. and Chandra, V. Alphanet: Improved training of supernets with alpha- divergence. In International Conference on Machine Learning, Pp. 10760-10771. PMLR, 2021.
Xiao, J., Hays, J., Ehinger, K. A., Oliva, A., and Torralba, A. Sun database: Large-scale scene recognition from abbey to zoo. In 2010 IEEE computer society conference on computer vision and paltern recognition, pp. 3485-3492. IEEE, 2010.
Wang, W., Bao, H., Huang, S., Dong, I.., and Wei, F. Minilmv2: Multi-head self-attention relation distillation for compressing pretrained transformers. arXiv preprint arXiv:2012. 15828, 2020a.
Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Gao, C., Huang, C., Lv,, C., et al. Qwen3 technical report. arXiv preprint arXiv:2505.09330
2025.
Wang, W., Wei, F., Dong, L.,, Bao, H., Yang, N., and Zhou, M. Minilm: Deep self-attention distillation for agnostgc compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33: 5776-5788, 2020b.
Yang, Z., Xu, Q., Bao. S., Cao, X., and Huang, Q. Learning with multiclass auc: Theory and algorithms. IEEE Trans- actions on Pattern Analysis and Machine Intelligence, 44 (11):7747-7763, 2021.
12

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-B--Dvergence
Yang, Z., Xu, Q.. Bao, S., Wen, P., He, Y., Cao, X., and Huang, Q. Auc-oriented domain adaptation: From theory to algorithm. IEEE Transactions Pattern Analysis and Machine Intelligence, 45(12): (111)):14161144174 2023a.
Yang, Z., Xu, Q., Hou, W., Bao, S., He, Y., Cao, X., and Huang, Q- Revisiting auc-oriented adversarial training with loss-agnostic perturbations. IEEE Transactions Pattern Analysis and Machine Intelligence, 45(12):15494- 15511, 1, 2023b.
Yang, Z., Xu, Q., Wang, Z., Li, S., Han, B., Bao, S., Cao, X., and Huang, Q. Harnessing hierarchical label distribution variations in test agnostic long-tail recognition. arXiv preprint arXiv:2405.07788 2024.
Zagoruyko, S. Wide residual networks. arXiv preprint arXiv:1605.07146. 2016.
Zagoruyko, S. and Komodakis, N. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. arXiv preprint arXiv: 2016.
Zhao, B., Cui, Q. Song, R., Qiu, Y., and Liang, De- coupled knowledge distillation. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 1 1153--1966,, 2022.
Zheng, K. and Yang, E.-H. Knowledge distillation based on transformed teacher matching. arXiv preprint arXiv:2402.111 11.4002.14444 148, 2024.
Zhou, H., Song, L., Chen, J., Zhou, Y., Wang, G., Yuan, J., and Zhang, Q. Rethinking soft labels for knowledge distillation: bias-variance tradeoff perspective. arXiv preprint 1rXiv:2102.00655 2021.
Zhou, K., Yang, J., Loy, C. C., and Liu, Z. Conditional prompt learning for langunlan models. In Proceed- ings of the IEFF/CVF conference on computer vision and pattern recognition, pp. 16816-16825, 2022.
Zhou, X., Ye, W., Lee, Z., Zou, L., and Zhang, S. Valuing training data via causal inference for in-context learning. IEEE Trunsactions on Knowledge and Data Engineering, 2025a.
Zhou, X., Zhang, M.. Lee, Z., Ye, W., and Zhang, S. Hademif: Hallucination detection and mitigation in large language models. In The Thirteenth International Con- ference on Learning Representations, 2025b.
Zhu, Y., Lu, S., Zheng, L., Guo, J., Zhang, W., Wang, J., and Yu, Y. Texygen: A benchmarking platform for text generation models. In The 41st international ACM SIGIR conference on research & development in information retrieval, pp. 1097-1100, 2018.
13

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via e-B--Dvergence
Appendix
Table of Contents
Prior Arts
16
Continuous Extension of a-B-Divergence
16
More Discussion Related Tracking Probability Allocation with Log Mass Ratio
17
C.I The Relationship between Log Mass Ratio and Logit Gradient
17
C.2 The Relationship Between Overall Gradient and Logit Gradient
17
a~3-divergence: Further Analysis on Trading off Hardness-Conneentaaion and Confidence-CCnceettatiio
17
Comparison of Gradients: FKLD, RKLD, and Our Framework
18
Pursuing Proper Mass allocation via a-Divergence
19
G Proofs
21
G.1 Proof of Proposition 3.1
21
G.2 Proof of Theorem 3.2
21
G.3 Proof of Proposition F.2
23
G.4 Proof of Theorem
25
G.5 Proof of Proposition 4.2
33
G.6 Proof of Theorem D.I
35
H Algorithm Protocol
36
Additional Experiment Settings
36
LI
Natural Language Processing Tasks
36
L1111
Datasets
36
1.1.2
Competitors....
37
I.1.3
Implementation Details
37
I.2
Vision Tasks
38
1.2.1
Datasets
38
1.2.2
Competitons
39
I.2.3
Implementation Details
39
J Additional Experiment Analysis
40
J.1
Natural Language Processing Tasks
40
14

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3-Divergence
J.1.1
Effects of SGOs
40
J.1.2
Distilling from Stronger Teacher
41
J.1.3
Comparison with More Baselines
41
J.1.4
LLaMA Family Distillation
42
J.1.5
Qualitative Evaluation
42
J.2
Vision Tasks
43
J.2.1
Distilling from Stronger Teacher
43
J.2.2
Cross-Architecture Distillation
43
J.2.3
How does ABKD perform with alpha/beta outside [0,1]?
44
15

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via c--3-
A. Prior Arts
Knowledge distillation (Hinton, 2015) is a promising technology for transferring knowledge between different models. The typical setup assumes the presence of a larger teacher model with more parameters and smaller student model with fewer parameters. To achieve this knowledge transfer, common approach is to let the student distribution qe mimic the teacher distribution p by minimizing a distributional measure D(pl(ge). This approach is referred to logit-based distillation (the focus of this work). Another promising approach is to leverage the rich information in the intermediate layers of the model, such as the attention matrix (Zagoruyko & Komodakis, 2016; Sun et al., 2020; Jiao et al., 2019; Wang et al., 2020b;a), the embedding features and their relationships (Romero et al., 2014; Liang et al., 2023b; Sun et al., 2019; Liang et al.. 2023a; Lv et al., 2024; Tian et al., 2019), etc. These methods are known as feature- based distillation. Due to the success of KD, it often outperforms supervised fine-tuning and has led to improved performance in various downstream tasks, including image classification (Kim et al., 2024; Yang et al., 2021; 2023a; 2024), instruction generation (Gu et al., 2024b; Zhou et al., 2025b;a), neural architecture search (Wang et al., 2021), and object detection (Li et al., 2024; Lv et al., 2024).
Logit-based methods aim to minimize the distance between student and teacher distributions, which have achieved profound success in the past few years. To do this, one can choose different distillation objectives, such as Maximum Mean Discrepancy (Huang & Wang, 2017), Total Variation Distance (Wen et al., 2023), Wasserstein Distance (Lv et al., 2024), or Pearson correlation coefficient (Huang et al., 2022). Most prior methods (Hinton, 2015) use primarily forward Kullback-Leibler divergence (FKLD) to let the student distribution to mimic the teacher distribution. On this basis, variety of methods have been proposed help the student learn better from the teacher distribution, such as using asymmetric temperature scaling (Li et al., 2022), decomposing the teacher distribution into separate learning of target class and non hon-target class knowledge (Zhao et al., 2022), removing temperature scaling on the student side (Zheng & Yang, 2024), normalizing logits (Sun et al., 2024), reusing the teacher's 's classifier (Chen et al., 2022), utilizing inter-class relationships (Lv et al., 2024; Jin et al., 2023), and so on. Despite achieving profound success, recent research points out that due to its asymmetry. FKLD tends to cover the entire support of the teacher's s distribution, leading to an overly smoothed student distribution. To address this issue, many works resort to using reverse Kullback- Leibllr divergence (RKLD) (Lee et al., 2023; Gu et al., 2024a; Kim et al., 2024; Gu et al., 2024b), which forces the student distribution to focus on a few modes in the teacher's distribution. At the same time, some works explore more general distribution measures (Wen et al., 2023; Agarwal et al., 2024; Wang et al., 2021) and composite metrics (Wu et al., 2024; Amara et al., 2022; Binici et al., 2022). Recently, some studies (Ko et al., 2024; Wu et al., 2024: Wen et al., 2023) have found that the superiority of FKLD and RKLD depends on the task and dataset. However, systematic studies providing theoretical insights into the suboptimal performance of FKLD and RKLD are either scarce or predominantly qualitative. This limits further exploration in this field.
this paper, we propose a generic distillation framework based on a-ve-divergence. Unlike previous generic methods, our approach is 1) built on balancing the unified space of our framework, we 2) theoretically explain why FKLD and RKLD lead to suboptimal performance, which further complements previous empirical observations. Fortunately, our framework 3) allows for flexible interpolation between them, ensuring better performance. Furthermore, we 4) confirm the effectiveness of the newly introduced distillation objective through extensive empirical experiments on language and vision datasets.
Contributions:
concentration and confidence-comcentratioo Based on analysis in
B. Continuous Extension of a-33-Dieereence
The u-B-divergence can be extended through continuous extension (by applying L'Hopital's Rule) cover all values of a, B. E R. Its more explicit form is defined as follows.
Σ. - 1 (()) a - 8
a+3 -
for a, R,B,a + 3 # 0,
as
Σx a p(())a (Inp(k)- - -nq(())))
1
- =(k)= + (.b))
for
a a 0,3 0,
(o...) AB
((() -1
Ex a -
p(k] -
for
a = - 3 0, a 0,3 + 0, for o, B = 0.
(8)
Ea t ( (k)) +kk)))
for
Σ 2 (hp(k))-ng((k)?
16

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3-Divergence
C. More Discussion Related to Tracking Probability Allocation with Log Mass Ratio
C.1. The Relationship between Log Mass Ratio and Logit Gradient
For given divergence algorithm A (p g). consider using the gradient descent method to update the loss function ! w.t the logits fy. then the distribution at the next step gf+ I is given by:
exp(Jx+1)
(1)
Σn exp(Jx=1)
exp(S; - s;
(9)
En exp(s - (2
exp(- 13
= =v(()) Σ (k) expl - v
Observing that the denominator serves as a normalization constant, this can be rewritten as:
(41(2) exp (-7V5;²)
(10)
(v)
Taking the logarithm both sides, get:
q (g) = 1 v/ + N2),
log
(11)
Lt (g)
where (y) denotes constant normalization factors independent of y. This indicates that the log mass ratio is proportional to V F.
C.2. The Relationship Between Overall Gradient and Logit Gradient
We first give the relationship between overall gradient and logit gradient:
vw yll
(12)
In this case, J is the Jacobian matrix representing the gradient of logits w.r.f. model parameters, and its dimensions C'x M, where C' is the dimensionality of logits, and M is the dimensionality of the model parameters W. Typically, we have
M >> C. For example, for an image classification task using ResNet-110 10 on CIFAR-100, C 100 and M 1, 110, 240. In the case of instruction generation tasks, for GPT-2 XLarge, C = 50, 257 and M = 1, 500, 000, 000. Thus, the matrix J is close to being full rank C.
In this case, w + 0, then we must have:
e +
(13)
Since the Jacobian matrix J is full rank, the product can only approach zero V + O, i.e. v + for all class channels y.
D. a-33--a diivergee::: Further Analysis on Trading off Clrneness and Confeeenee-onnenttatiion
Of course, there is no free lunch. broad hyperparameter space offers more flexibility but also makes finding suitable values more difficult. While grid search could theoretically yield optimal performance, it introduces additional overhead in framework. Fortunately, guided by the following theoretical insights, we can design principled divergence algorithm for the target task with inductive bias more efficiently.
(o. Theorem D.I. Let de41 (y) be the distribution obtained after one gradient step, starting from qe using the a---divergeece. Define A () as the difference of log mass ratios across two classes and y2, obtained from the a- o-B--dieergece::
(a.9) (v1. V2) & LogR! (V1) LogR( (y2). .y))).
(14)
17

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergen
We have the following (for appropriate posilive constants S, ò2 82, and any real numbers on and og in the range (0, 1] satisfying on oz)):
1. a- a-ß-divergence transfers probability mass from overestimated classes to underestimated classes more aggressively as decreases. If y and y/2 are such that ô1 < q:(1)) - (v(22) S p(y)) (where 81 > 0), and p(y)) 2 p(y2) + G. holds that
12) A (02.9) (1))4)).
2. a- a-B--ivergence reduces the probability mass of classes with larger enror (p(y) - ge(') more aggressively as a decreases. If y1 and y/ are such that p(y1) < (((1)) = qe(12) S 62 (where 82 > 0), and p(y)) 2 p(u2) + G.
holds that A .9) (yn, v2) 2 A (0 (y1, 92)-
3. The a-B-divergence becomes more (less) preferential in focusing the error on classes with higher student confidence as 3 increases (decreases) when reducing LogR
The proof is in App. G.6. Case shows that a smaller a (stronger hardness (rndness-cooncentraton)) leads to aggressive mass reallocation across classes when some classes are overestimated. Case 2 shows that a smaller a will more aggressively penalize overestimated classes with large errors (a.k.a,, hard classes). this sense, a proper assignment of a leads to a better harddess--concentraiio effect. On the other hand, case 3 shows that a larger /3 tends to focus more on reducing the error from classes with high student confidence and ignores the error from other classes, shown in Fig. 6(c) and Fig. 6(d), resulting in sharper student distribution. As such, one select proper 3 to ensure better effect.
plperparameeer tuning guidelines. In principle, one may prefer to select a larger 3 than O, which should be inversely proportional to the ratio of non-target classes in the output distribution. A proper 3 allows the student to effectively learn from the teacher's soft labels while maintaining an adequate focus on the target class. On the other hand, choosing a smaller than leads to more aggressive probability mass reallocation across classes. Therefore, small and proper a can more effectively avoid local optima. This becomes particularly important when the distributions are far apart, shown in Fig. 6(b).
Empirically, we find that for tasks with low- www-dmmensioall output distributions, such as image classification on CIFAR-100, selecting a large and small 3 is sufficient to achieve optimal performance, as shown in Tab. and Tab. 4. However, for tasks with more high-dimensional output distributions, such as instruction generation on the Dolly dataset, selecting a small (which leads to more aggressive reallocation of probability mass) and a large 3 (to emphasize learning the soft label information) are crucial for achieving exceptional performance, as shown in Fig. (aa) and App. I.1.3.
E. Comparison of Gradients: FKLD, RKLD, and Our Framework
Lemma E.1 (Cichocki et al. 2011). Given two distributions and the gradient of the a- a-B-divergence with respect to 0 is calculated as:
(). (p (0)
- -
Ogo
[k)
-0
p.go
(15)
00
0e
weights
a--ooom
where pp.s is the ratio between arbitrary distributions and
Proof. The formula for the -----divereence is defined as follows:
IX ((a..)) (p go)
a, Σ p(k)) (0)) (k)) a + g p(k))tt
a + 40 (k))+
(16)
Taking the derivative of each term with respect to the parameter 8 we have:
OD (a.9)
AB
(p
|| go)
Σ
(a.3) (p qo)
0qo (k)
(17)
20
(kk(k)
09
h
18

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via c- a-3-Divergence
where
i
D D (p g0)
1
(8 Bo0 (k))at(-1)
dqe (k)
a3
(o(k))(((k))8-)) 1000--0010 (k0
(18)
p(k)
9o(k)
-
a
Substituting Eq. 18 into Eq. 17, obtain:
p(k) a (1)
aD (a.
qo)
(k)(k)
Σ
-0(k))+t
(19)
00
00
This concludes the proof.
The gradient of the FKLD with respect to is given by
a 30
(qo(k)
p(k)
DKL (pllga)
Σ
(20)
00 qa (k:
The result is the negative gradient of the model probability, inversely weighted by its value. As Ko et al. (2024) stated, when qo (k) = 0, the gradient norm increases, causing large, noisy updates that can hinder optimization. Similarly, the derivative of the reverse KL divergence with respect to 0 is given by:
i
Sgo (k).
(log (o(k) p(k) +1)
Te (gollp)
Σ
(21)
08
00
This value becomes very large when p(k) O. our framework, by Lem. E.I, the derivative of the a-B-divergence with respect is:
p(k) (o(k)
aD (p go)
Ōq(k) 00 (k))aat--11
00
Ck
When 1 and 8 0, FKLD becomes special case; when 0 and B 1. RKLD becomes special case. The parameter controls the focus large or small ratios p/ge. while 3 adjusts the weighting of these ratios through the scaling factor 4g - When choosing
aD () (pllge) tends to excessively focus on the extreme values of p/q (i.e., when 00
p(z) and q(x) 0). On the other hand, choosing close to 0 would overly focus the extreme values of q/p (i.e., would place more focus p/g of classes with high student confidence, while conversely treating all classes more equally. This means they provide a fine-grained way to tune the model to emphasize specific likelihood ratio ranges, thereby ensuring stable gradient optimization.
when (=x) > ( and p(x) ss 0). Additionally, choosing larger value of a -
F. Pursuing a Proper Mass allocation via a-Divergence
To achieve an effective trade-off between hardness-conceentatioon and confidence-conccentraiioo can, for example, extend FKLD and RKLD to family of generalized divergences by introducing additional dimension a, which is known the a-divergence (Chernoff, 1952).
Definition F.1 (a-divergence). Consider a IR 10, the a-divergence two distributions is given by:
Da 9) A 1 ((()))(
ala 1)
where p [p()))]=1 and q = [m( are two discrete distributions over C classes.
19

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-B-Divergence
Remark. The a-divergence includes (p qo) as a + 1, and (gx p) as a - 0.
The following proposition characterizes the effect of hyperparameter on reducing [LogRf The proof is in App. G.3.
Proposition F.2. The updates induced by a-divergence for qe within one gradient descent step are given by:
ply) a (2))))
p(k)
(k())
(o) (w) S n ( 1-o
+g.(v) Σ (k())
+ +1Ng
a
(a)
(a)
(b)
(6)
where Nº (y) denotes constant normalization factors.
This proposition indicates that term (a) scales the relative importance of large versus small 4u()), controlling the effect. Term (b) adjusts the relative emphasis of the error between p(y) and q (v), control- ling the hardness-concentrann effect. Together, these terms are coupled, with their interaction governed by a and - a. Unfortunately, this unnecessary constraint makes it intractable to adjust their effects independently.
The following theorem validates our idea and shows that the a-divergence can only inflexibly interpolate between FKLD and RKLD (Fig. le) in a linear subspace of the planar space formed by terms (a) and (b), shown in Fig. (()) The proof is deferred to App. G.4.
Theorem F.3. Let (3) be the distribution obtained after one gradient step, starting from qu using the a-divergence. Define A as the difference of log mass ratios across two classes y1 and y2. obtained from the u-divergence:
(11, 82) e LogR (v1) - (g2)(((22)
We observe the following linear trend (for appropriate positive constants Ç. ō1, 82 and any real numbers on and 0r in the range [0, 1] satisfying 01 < 02)):
1. The a-divergence transfers the probability mass of overestimated classes to underestimated ones more aggressively as a decreases. If y1 and y2 such that p(yz) 81 < 90((1) = 9: (12) ply1) (where 81 > 0), and ply1) 2 p(y2) + +. it holds that A (y1, 122 2 - (yr, V2).
2. a-divergence reduces the probability mass of classes with larger error Ip(y) - qe()) more aggressively as a decreases.
If yı and y2 are such that p(y1) < qe(y)) (v22 S 1 62 (where 82 > 0), and p(y1) 2 ply2) Ç, it holds that
A (y1, V2)
(21,
3. a-divergence increuses the probability mass more preferentially on underestimated classes with larger probabilities qe (y)) as a decreases. yn and y/ are such that qe (y2) + S q+(in) E - 82, and p(yn) = p(yz) > co .4(y)), where co a positive constant 1, it holds that A (y1, Z A 42).
4. a- a-divergence reduces the probability mass on overestimated classes with larger probabilities qe(v) more conservatively as a decreases. If yı and y2 are such that qe(y2) + S. ge(y1) S 1 1-1100 82. and co qr (y2) < p(y1) = p(y2) < ci (. where en and ei are constants with eo > and C₁ < 1, it holds that A (2) Z (y1,
These cases illustrate the differences in probability mass allocation for different Specifically, from Case I and Case 2, it can be seen that smaller leads to more aggressive reduction of the probability mass on overestimated classes, transferring it to underestimated classes. Case 3 and Case 4 show that smaller a (or larger 1 a) tends to concentrate the probability mass more on classes with high student confidence.
In summary, we have the following conclusion.
The a-divergence achieves suboptimal balance between hardness corddness-comccetrration and confidence- inflexibly.
20

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via Divvereeeeeeeeee
G. Proofs
G.I. Proof of Proposition 3.1
Lemma G.1 (Tajwar et al., 2024). For a given distribution qt, the algebraic relationships between LogR and the gradient of the logit fy with a given learning rate y in FKLD and RKLD are given by:
FKLD:
LogR (y) = 17- (p(y) - qu(v)) + (1),
RKLD: LogR (y) - 11 (e(v) (logply) - (e + Σ gr (k) log (k) (oggolog()))) + N,
where N. (y) and N (y) denote constant normalization factors independent of y and vanish to zero when qt-
Restate of Proposition 3.1. The updates induced by FKLD and RKLD for q: within one gradient descent step are given by:
FKLD:
|LogR (v) S n 1 (p()) (v( + N₂ (v)
RKLD:
LogR R. (v)l S n (v(
(oggpl) - log ge (v)
+ Σ (k)
log p(k) - log ge (K) + (w)l:
(01)
(b1)
(02)
(b2)
where (1)
(v) and NR(2)) (y) denote constant normalization factors independent of y and vanish to wero when q
Proof. By Lemma G.I, and applying the triangle inequality, we can directly obtain the following bounds:
|LogR (v)l s 1- pp()) q.(x)] + NN (v)l
(22)
[ogR] (v)l S 17 u:(v) (ogg(y) - log dr (v)l + Σ de log)))) log (k)) + |N, (w)),
(23)
This completes the proof.
G.2. Proof of Theorem 3.2
Restate of Theorem 3.2. Let 1 (v) be the distribution obtained after one gradient step, starting from using the FKLD. Likewise, let 14+1(0 (w) be the distribution obtained using the RKLD,, from qt- Define of and A as the difference of log mass ratios across two classes y1 and y2, obtained from the forward and reverse divergences respectively:
Df (v1, A LogR," (v1) - LogR, (12).
and Di is similarly defined. Then we have the following (for appropriate positive constants S, 8, 82):
1. RKLD transfers probability mass from overestimated classes underestimated classes more aggressively than FKLD. If yı and Y2 are such that p(y2) < 81 < <((()) (322 < p(v1) (where 8N > 8, > 0), but p(v1) 2 p(y2) + S.
then, (11,, > D (11.122) W2).
2. RKLD reduces the probability mass of overestimated classes with higher error (p(y) 4<(y)) more aggressively than FKLD. If yı and y2 are such that p(y1) qo(11) (12) 82. but p(y1) 2 p(y2) + C. then, (1 U2)
D (u).
3. RKLD more preferentially increases probability mass on underestimated classes with larger probability qi qv(v) than FKLD. If yn and y/2 are such that qe (v2)) + C S 4e(1n) < - 82, and p(y1) = p(y.) > co ge(y)), where Co is positive consiant > 1, then, AT u2) a (2).
4. RKLD reduces probability mass on overestimated classes with larger probability qe(y) more conservatively than FKLD. If yı and y2 such that qi (122) + I (1(11) < 1 82. and co qo (2(2) < p(v)) p(y2) c₁ qe(vr) where co and Or are constants with Co > and 61 < 1, then, A (v1, y2) A! (1),
21

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via Diveeggeeeee
Proof. For Case note that (vh1 (v2) < p(v1), based Lem. G.I, have:
A (y1 U2) 77 (p(in) p(v2))
(24)
A (11 v/2 = 77(vv) (log logr(v2))
(25)
The discrepancy between A and A' is now given by:
(y1, 322 - (vv)) logp(v)) - log p(y2) - p(v1) p(v2)
A (11,, y2)
(26)
(
Notice that by the lagrange's mean value theorem, there exists a co E (p(v2), p(y1)) such that:
dlogg
log p(y1) log p(22)
(p(v)) p(v2)).
(27)
dp
p=co
Since dlogg dp
1
we have that:
pmec
(((i)
(y v2) - A (11182) 1/2 = 17 (p(vi) - p((2))
(28)
This quantity is positive if ((v)) > co = 81. Then, for Case similarly, since p(y)) qe(v)) 2o (w2). have:
(g(nn)
(v v2) (v),,
(p(v1) - p(12))
(29)
co
where (p(v2). p(y))) )) is obtained by applying the lagrange's 's mean value theorem to the difference log p(y)) - log p(v2).
Notice that ((1) > p(y1). ). so this term is always positive.
The proof for Case consistent with Tajwar al. (2024), and is omitted here.
Finally, prove Case 4. Noting that ((22) < p(v1) p(v2) < 4(v1), we have
A (y1, 32) n(1(vz) q(y))).
(30)
Additionally, for (311 32). we have
A (111 v2) = nia(vn) - ((22) log p(v)) nla(v1) log(g(v1) - ((22) log(1((>))
(b)
(31)
+ m(((())
20
For item (b). by the lagrange's mean value theorem, there exists a point co E (4(12), ( ) such that
dg lglog
(1())) ((v2) log(a(v2)]
(4(12) q(v2))
dq
(32)
q=co
= + log (a(3n) ((2)))
Substituting into Eq. 31, we obtain
A (y1 U22 n(g(v1) - ((v2)]llog p(y1) log co 1 DeL (pll)))
(33)
We need prove that
n((())) ((1e)) - log co (pllg) - q(22)]
(34)
It suffices to choose sufficiently large p(v1) (q(32). q(v1)) to ensure that
p(y)( log + 1Dxx (pllg) > 0.
(35)
This condition can satisfied.
22

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--Divergence
Remark. These cases highlight the contrasting behaviors of FKLD and RKLD in certain scenarios. Specifically, Case I shows that RKLD is more aggressive in transferring probability mass from overestimated classes (i.e., p(y) < ((v)) to underestimated ones (i.e., ply) g)y)). Case 2 indicates that when two classes have the same predicted values g(v) greater than the target values p(y). RKLD more aggressively penalizes hard classes with larger overestimations (i.e., larger error (p(v) (.))). Case 3 indicates that when the predicted values g(y) for two classes y1 and y/2 are both below the target value ply), RKLD preferentially increases the probability mass the class with larger student confidence g(y), even though both predictions are equal at the target value p(y). Case suggests that when one class yi overestimated and another class y/ underestimated, RKLD reduces the overestimation of eigh-probability classes more conservatively.
In summary, combining Case 1 and Case 2, conclude that RKLD reallocates the probability mass across classes more aggressively by penalizing errors in hard classes compared to FKLD, aiming to achieve better matching. This can help avoid local optima more quickly in some cases, ensuring superior performance and faster convergence throughout the training process, as shown in Fig. 2. From Case 3 and Case 4, it can be inferred that RKLD tends to concentrate the probability mass on a few classes with high prediction confidence. contrast, FKLD reallocates the probability mass evenly across all classes since the weights for different classes are identical. This can lead to sharper distribution when using RKL.D, as empirically validated in Fig. 6(c) and (d) (B 0 for FKLD and for RKLD).
G.3. Proof of Proposition F.2
Lemma G.2. The gradient of the softmax function ((1) = Σ1 cf with respect to f3 is given by:
Og(1)
(()((1 a(1))
ii==j. j. j.
(((()((5)
Proof. To derive an(i) 0fz we consider two cases: i i= and i
1. Case i Using the quotient rule, we have:
0q(1)
eh Exehe -ch
0f:
Cxeedey
(36)
ch cl ch.
Σn ehu etep"
Simplifying:
Og(i)
q(i)(1 g(i)).
(37)
Of,
Case i * j: We have
0g(1)
cticls
Exeraja
(38)
=-g(i)(í).
In conclusion, the derivative is:
0g(1) 05, (-q(1)((1) i.
(1()(11- - g(i)) ifi=j,
(39)
Restate of Proposition F.2. The updates induced by a-divergence for q: within one gradient descent step are given by:
log (x) (x) I ( I-a
p(y) (
p(k)
+q((v) Σ (k) 1-a
+ +(N7())1
a
a
(4)
(4)
(b)
(6)
where No (y) denotes constant normalization factors.
23

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via c- -ivvergDee
Proof. The formula for the a-divergence is:
Da(p g) _
1
Σ p(K))
(40)
ala - 1)
Using the chain rule, we have:
0
aD (pl q) 0q(k)
Da 11 -
(41)
0f,
oq(k)
where
aD (p 9)
ply)
(42)
(3
a (1y)
and
04(y) 0 eiv
(43)
0f, 0f, Σx etx
Combining Lem. G.2, the Eq. 41 can be expressed as
0
(D, |l g) (q(k)
aD. (p 11 9) 0g(y)
(D
②)
Σ
05
0g(k)
Of(y)
0q(y) Of,
kzy
)) (-q(k)g(()) + (-1 (P(U) )) (0(1)(1- -(1)))) (((k)
p(k)
Σ
((())
(44)
kyl
Σ [a(k))->00((1 (((()((((()(()))))) p(())))
q(v)² (p(v)* - 9(v)
Now, consider using the gradient descent method to update the loss function with respect to the logits fy. then the distribution at the next step 4*+1 is given by:
(++(1)) = exp(J(+1)
Exex((Fh-1
exp(S] - -1 f3
(45)
Exexp(s 1/ st
exp(-17,,*1) ,
=v Σe (k) exp(-
Now, substituting the gradient formula of the a-divergence, the characterization of (f+1(())) is obtained as: exp ( [a(y)>=a (p(y) (()))) + + ((v) (k))) p())))))) (i) exp (2 (a(1)2=a (p(i)a (()))) + ((1) <x
(() qQ)))
(46)
((k)
Observing that the denominator serves as a normalization constant, this can be rewritten as:
d (v)
! a(v))
- + g(v) Σ a(t) ((())))(((())))))))))))))))))( -
Ze (v) exp
(47)
Taking the logarithm on both sides, we get:
-g(x)) + (0()) Σ (K)))- Pi (v) = 2-( ((())
(ge(k)² p(())
log
(v).
(z)
a
a
(48)
where NO (y) denotes constant normalization factors. We can further derive that (x)) 1-a
log (1)
p(y) 4
p(k) (k)a
+q((v) Σ
q (k)
(49)
(v(v)
a
(4)
k
(a)
(b)
(b)
This completes the proof.
24

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Dierggen
G.4. Proof of Theorem F.3
Lemma G.3. Let p and q be normalized probability distributions over class numbers C' ie., Σe p(k) EEp(k)
Et 9(k) =
Define the function F F(a) for a [0, 1] as:
F(a) a 1 2 2 1- (1-Ep(k)a(k))
Then, F(o)) decreases monotonically as a increases, and F(a) 2 0 on [].
Proof. Define:
5(0) - Σ p(k))
(50)
Thus, the function F(a) can be written as:
F(a) 1 -(())
(51)
a
First, compute the first derivative of S(a):
(aa)) Σ (())) (Inp(k) ((np(k)--ua ln))
(52)
k
and the second derivative:
(a)a) = Σ p(k) [lk) p(k))'((k))* (())'=aa (Imp(k) - In )(k))
(53)
2 0.
Since (a) 2 0 for all 8(a) is convex function of a.
Now, compute the derivative of F(a)::
F( d
1-1- S(a)
da a
-S a - (1 S(a))
(54)
a²
N(a) a²
where
((a) e (a) 1+ 1+ S(a).
(55)
To show (a) S 0, it suffices to prove ((a) S 0.
Since S(a) is convex, ('(a) is non-dccreasing. Additionally, considering the boundary conditions:
N(0 = 0, N(1) = -DL (pllg) 00
(56)
and since
(' d (S(a) aS'(a) - (a)
da
(57)
0,
((a) is monotonically decreasing. Therefore, for all [0, 1]:
N(a) S.
(58)
25

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a--3-Divergenc
Thus,
N(a)
F((a) = a
(59)
I
Hence, F(a) is monotonically decreasing with respect to a on the interval [0, 1].
Finally, note that F(1) 0, and therefore F(a) 2 0. This completes the proof.
Lemma G.A. Consider the function f(a) _ p(x) (vs) ) where p(x)) and q(tk) are constants. The derivative of the
logarithm of f(a) with respect to is:
d
a(ve) p(z)
In
2(ve) p(x)
In (p(x)) a(v:))) Imp(x) +
(60)
da
9(ar)) p(z)
Proof. Start with:
In -(x))) pnp(x) + 11n11-11 (I(us)) p(x)) )
(61)
Differentiating with respect to a, we get:
d
In (p(x)* ((v))))= d d
lmp(x)) In (1 (1- ((vi) a p((x)
(62)
da
da da
The derivative of ln is Inp(x), and the derivative the second term is:
g((ua) p(z)
In (v,)
g(U)) p(x) ") =
p(z)
(63)
da
2(0s2 p(z)
1
Combining these gives the desired result:
2(30
In o(vr)
d
In (p(x)* -(9)))) inp(()) +
p(x)
p(x)
(64)
da
1
((ars)
Lemma G.5. Let 1 2 a 0 and define
h(s) 2 1 - 20 2as In
(65)
for 0 < < 1. Then h(ss) is strictly decreasing on (0, 1).
Proof. Since < s 1, we have lns < 0, thus In 8) = In Substitute this:
h(s) 1 - 22 2as Ins.
(66)
Differentiating term-by-term,
N((s) = 20s
2a-1 2a + -2a((us a-1 lus + sa-1)
(67)
= 2 a- la lolns+1-89)
Set
((s) = anns+1s55*
(68)
As + 0+, Ins -000 and thus q(8) -00 At (1) -a 0+ 0. Moreover,
d((s) a(1 - 5a)
(69)
> 0.
26

---

0, ög ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergence
Since 0 8 < implies 5th 0. Hence g(s) is strictly increasing on (0, 1), and have q( 0. Thus g(s) 0 for all 0 < 1. Since 2as > 0 and ((s) < O, we have (s) < 0 for all 0 < < 1. Therefore, h(s) is strictly decreasing on (0, 1).
Lemma G.6. For 0 < < 1 and 1 2 O, define
B(s a)) - 2(1 (1 + 8*) In (1+ s
(70)
Then ((s, a) is strictly increasing in s on (0,1) 1) and increases from
as goes from
Proof. Since < 1, have In and thus
In s] In s. Substituting this into Eq. 70 and simplifying, we obtain:
2(11 8*)
((s,a) (1+ + 8a) Ins
(71)
To differentiate 3(s, a) with respect to s, let:
f(s) 2 1 8 ((()) = -os -as a-
(72)
g(s) 2 11+5 8", g((s) = as
(73)
h(s) 2 Ins|| = lns, 11(s)
(74)
Thus,
3(s,a) 2f(s)
(75)
g(8)h(8)
Applying the quotient rule:
d3 (())(()))) (s))) ((S)g(s)((ss)
(76)
ds
[g(s)h(s)}²
Plugging Eq. 72, Eq. 73 and Eq. 74 into Eq. 76 and simplifying, we arrive at:
dB
2h(s) ds (1 + 2ln
with M()) = 2a 22xs In
(77)
From Lem. G.5, it follows that h(s) decreasing (0, 1). As 1- h(s) + Therefore, h(s) on (0,
Since (1 +++s" s0) P > 0, it follows from Eq. 77 that d3 0. Therefore, ((s, is strictly increasing. Furthermore, taking limits:
lim as) 0, lim B(s, a) =a.
so ((s. a)) increases from 0 to as goes from to 1.
Lemma G.7. Let f(s, a) A s" s))
(1- E where 0 < s < and a > 0. Then, f(s, ai) is strictly increasing with respect to for
all 0 < < 1 and > 0, i.e.,
8f > O.
ds
Proof. To compute
we use the quotient rule. Define:
- a (lns)2, v 21- f(s,a) a =
u
(78)
The derivative is given by:
0f
ul un!
(79)
Øs 2
27

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-
First, compute w/:
u == s d' al=as=-((nns) + 20--1111 Ins = 5" [a(In s)² 2 + +22 11
(80)
Next, compute v':
= (1 - (1- 2(1 (-ass")) -2ass" (2 (1-s"))
(81)
Substituting u' and v' into the quotient rule:
0f 8" [) + 2n (1 (1= 2as 20
(In)) 8n)
(82)
0s
(1 8a)a
Simplify the numerator:
Numerator = -0- 5n) s) + 2n + -ass2a-((nn)) (lns)2
(83)
Factorize:
Numerator = 80-1 - 54) [a(1 + s (Ins)² + 2( - sn)
(84)
Thus:
0f 8a-1
[a(1 [a + 80) 8)2 + 2(1 - In)
(85)
as
(1 - 8a)3
To determine the sign of 0f ös note:
ga-1 since
and 0,
(1 - (1=5a)$ 0 since 0 and 0,
Ins < 0 for 0 < 8 < 1, hence (Ins)² 8}² > 0.
Denote the remaining expression inside the brackets as:
N N a(1 + +8)) + 2(1 5²) (P)nn
(86)
Rewrite N as:
N =
2(11 2(1-5a)
(87)
Since In > O, we analyze a(1 + + In 2(1 - 8*) > 0:
2(1 )
a >
(88)
(1 + 8d) Ins
From Lem. G.6, it follows that for all 0 < < and a > 0 > (ss a:) always holds, implying:
N> 0.
(89)
Hence:
05 > 0,
(90)
Øs
which shows that f(8, a:) is strictly increasing with respect to s.
Lemma G.8. Considering p(x), q(1)) [0, 1] such that p(æ) q(v1) 2 q(z2). and 0, the following function:
F(a)= ((11) 1-a p(x) 9(v1)a 4(12)1-a p(x:)a 4(122)
(91)
F(a)) decreases as a increases.
28

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via c--3 Divergene
Proof. Let h(a)) A In F(a), we get
h(a)= = (1 a1-a) In (v11 (12)
+ (p(x)* g(v1)*) In (p(x)* 9(v2)
(92)
Diffcrentiating h(a) with respect to we get
(a()) = - In g(v1) d (p(x)" - q(v))a) ((y2)
d (p(x)" - g(82)
(93)
do
do
Combining Lem. G.4, obtain the following expression for the derivative of h(a):
(o(o) - In ((1)
2(2 p(x)
In
2(vr) p(x)
g(sa) p(x)
In
p(r)
(94)
(1/2)
2(v1)
2(122) p(z)
a(x;) For convenience, define the variables s; p(z)
for = 1,2, where 1 12 2 $2.
Using this substitution, we obtain:
l((a) - In 81
8 In 81
$2 In 82
(95)
82
(a)
(b)
To analyze the sign (a))), consider the second derivative with respect to a:
d
( s In 81
$2 In 82
(a)a)) = da
+
(96)
1 1- 39
1 - $
Using the quotient rule, differentiate each term separately:
d 8 In
8 (In 8)
(97)
da
1- za
(1 5a
Therefore:
(a(a) = of 81)²
56 (In 82
+
(98)
(1 x2)2
)2)2
From Lem. G.7, we have
R" S 0.
Thus, have shown that (o)) is monotonically decreasing. Considering the limit of h'(a)) 0+, have:
lim l1(a) = a-0
Therefore, N' (a) 0 always holds, and F(a) decreases as a increases.
Lemma G.9. Define
f(p(V)) (log - log q(22))² 1112)1~a
+ (logp(ıı) - log 9(2)))² q(v))1' 1-a
where 0 < 4(v2) < p(v)) < 4(11). Then:
1. The function (p(())) is monotonically decreasing with respect to p(y1).
2. There exists a unique constant co E (g(yz), q(11)) such that f(co) = Moreover, for all p(y1) > co. holds that f(p(1))) < 0.
29

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-8--
Proof. The derivative of / with respect p(v1) is:
2
((((()))) = p(v1)
[logg(()) - (og((v))) ((1)) (logp(ın) - ((22) ((v2) [
(99)
Noting that 0 < ((v2) < p(v1) ((v)). therefore ("((())) - 0. When p(v1) Q3y2):
f(12) (log 92 log ( 1 1a-aa- 0.
(100)
When p(y1) 9(v1):
f(g)) -log q1 log Q2
92/2-)2000 a
(101)
According to the intermediate value theorem, since f(p) is continuous on p (92. q1) and decreases from a positive value to negative value, there exists unique co E (92, 91) such that f(co) Therefore, when p Cu. f(p) 0.
Restate Theorem F.3. Let (44(1) be the distribution obtained after one gradient step, starting from qt using the a-divergence. Define Ag as the difference of log mass ratios across two classes y1 and y2. obtained from the o-divergence:
(11 82) _ (og) - LogR((y)))
We observe the following linear trend (for appropriate positive constants Ç. 81, 82, and any real numbers 0i and 02 in the range [o, 1] satisfying an < 0z2):
1. The o-divergence transfers the probability mass of overestimated classes to underestimated ones more aggressively as a decreases. If y1 and y/ are such that p(y2) 51 qc())) (12) p(y1) (where 81 > 0), and ply1) 2 ((22)) +
it holds that A (y1, V2) 2 A (y1, 1/2).
2. a-divergence reduces the probability mass of classes with larger error Ip(u) [e()l more aggressively as a decreases. If yı and y2 are such that p(gn) < (c())) = (x(22) I 1 - 82 (where 82 > 0), and p(y)) 2 p(y2) Ç. it holds that
A (y1. 12) 2 0z (21.
3. a-divergence increases the probability mass more preferentially underestimated classes with larger probabilities ge (y) as decreases. If yı and 1/2 such that ge (322 + S ge(()) < 1 - 82, and p(y1) p(y2) > co -4((1),, where co is positive constant >> 1, it holds that A (y1, y2) 2 A (y1. (2).
4, a-divergence reduces the probubility mass on overestimated classes with larger probabilities (v(v) more conservatively as decreases. yı and y2 are such that q0(y2) + S (c(31) $ 1-11-00 82, and co qt (y2) p(y1) = p(92) c1 (6(1). where co and ej are constants with co and Ci < 1, it holds that A. (y1, V2)
(21,
Proof. First, we prove Case 1. Note that qe (3r) = qe (V2) = q(x), based on Eq. 48, we have
A = 1)((2)1 1-a p(y1)* )ar p(y2)
(102)
Let f(a)) = 1(x) #( plva)" and then it can be rewritten as
p(v1) 3 dy.
f(a) g(x)² -a
(103)
p(92)
Let t y q
substituting gives
p(yn]
f(x) = (0))² 2-a
aa-1 ((x))=-1_ ((x))d) di
p(ua]
(104)
w(x))
= (xx)
aa-1 dt.
4(c))
30

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via ae-3-Divergence
Using Leibniz's rule, get
sivll f'(a) g(Ø)
ta- In dt.
(105)
p plys
Note that the sign of (a) depends only on
In dt, so define
((2)
plunl 900)
h A
In & dt.
(106)
pinal
Clearly, the sign of h depends on the relative size of g(z) with respect to p(v1) and p(u2). To differentiate h with respect to ((x), using Leibniz's Rule, we get
Plu
hl()(x))) dg(x))
t
In itt
pland
(107)
p(v1)
p(12)
[ta-1 In
[
In
p(us) 00.3
q(x)33
p(us] DERT
q(>)²
Simplifying, we obtain:
N((g(x)) p(y2)* In
p(aal
p(21) 191 In
p(vi] ((x))
((2)
(x))at1
(108)
p(y2)" In p(v2) p(1n))()))))))))))) )a In p(v1) + (p(y)) - -P()2))) In
(x))11
Note that the sign of (g(x)) depends only on the sign of the numerator, and the numerator is monotonic increasing function of ((x) ((p(n)) 2 p(y2)). Note that ((x) S p(y1), we have
11(((x)) < 16 (v((1)))
pluz) In plya) p(y))
(109)
p(y) a++
5.0
Thus, we have proven that h is monotonically decreasing as g(z) increases. Also, since
h(p(22))
-- In dt
(110)
2 0,
and
h(p(y1)) =
8a
Intdt
M(u2) plaz)
(11I)
< 0.
By the intermediate value theorem, there exists E [p(v2), (11)) such that when q(x) > Ca. we have
h(g(x)) <0
(112)
for all values. By combining Eq. 105, we can conclude that f'(a) < O. Furthermore, let 61 2 max(ca) for any a E [0, 1].
Then, when p(y)) 2 q(x)= > 6,, we have fa( < 0 for a E [0, 1]. Thus, we have proven that for any aj and 0g E [0,
such that 0q < a2 we have
A (3/, 12) A (v1,
31

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via Diivergeeeeee
Next, prove Case 2. Similarly, since (111 (o(v2) g(x). deduce that
A = ng(x)==a p(y1) p(y2)
(113)
Let f(a) a g(z)==a. p(vx))-p((s)"
p(v)) get
pivil
f(a)= g(z)
ta dt.
(114)
pipa]
Using Leibniz's rule, we get
2(v1)
f'(a) = g(0)
q(()
aa- mndd.
(115)
m(u2)
v(x)
Note that q(x) > p(31). so
(a) SO.
(116)
Thus, this proves that for any on and 02 such that 01 < 02, have
A (y) 122 2 A, (v1. U2)
We now proceed to prove Case 3. Note that p(y)) p(y2) 2 ((v)) ((22) Ç. thus deduce
A (y), 92) mp(())' ' p(u1)
q(y1)"
(122)) a p(y1)" 9(v2))
a
a
(a)
(b)
(117)
n (q(1)) ((22))
Σ a(k))
(() 1-a
(g())) -((k))=<(())))))
(c)
First, through Lem. G.3, we know that (c) 2 0, and it increases as a decreases on [0, 1]. Then, we consider the term (a) (b): A (v 2 (a) (b)
(1((1)) p(y1) r(v1)"
p(y2) (322)
(1(12)2-a
a
(118)
- p (
(111)~~a p(u1 )a ((11) )a (122 1-a p(y2) ) (12)
From Lem. G.8, it can seen that term (2) increases as a decreases. Term (I) can be expressed as:
2- p(u2) q(12)a
p(2))
q(y2)==a
aa-1
dt
a
4(82)
(119)
n(ua))
7092
9(12)
to dt.
Noting that
p(v2) d
z(ua)
ul92)
go
dt
vis2)
aa- Intdt
da
(120)
2 0,
it follows that term (1) decreases as a decreases. Therefore, as a decreases, term (2) S 1, the value of (a) (b) increases
as a decreases, although its value remains overall less than 0. If term (2) 2 1, we have (a) - (b) 2 0, even though
a=1 (y1,y2) <
32

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--
Finally, prove Case 4. Since p(y11 p(v2) and g(v22 ((1),, we have
A U2) = (7(1))²~a p(81) - q(v1)" na(v))=aa p(B1)= - -3V2)
(121)
7 (g(v1) q(v)))
Σ (())t (a(k) - p({))))
Considering f(a) 2 (11,,,) 92)/n, taking the derivative with respect to a yields:
1 (((a) a²
(-1 a logp(v)) a log ()(()) ((()))
(1 a log) (logg(())))ppp
(a(v1) ((12))
1+ Σ - -kk)== (())a))
(122)
Σ
log ((k)) p(xx)* + (((x))))
+ *(k)) q(k))=a(( (l)))()))p(()) + log (()) ((k)a))
Noting that the sign of (a) depends solely on the numerator, let h(a) denote its numerator. Differentiating h(a) with respect a, obtain:
Na(a) ap(s))
p(v) - log + (logr(y1) log v(vi)))) (((()))))(v))(xx))
- (1(y)) - ((22)) Σ "pk))"""(k)""">) (log logr(k))
(123)
S + (logp(y)) log ((v))(())))))))))))))) (v11)))(2)))
p(y1) log ((z))))))) ((())
From Lem.G.9, it can be concluded that for any a, there exists a point ca E (g(x)), ((1)) such that when p(v)) > Ca. (a(a) < 0. Therefore, let 81 max(ca) for a [0, 1]. Then, when p(y1) > 81, h(a)) is monotonically decreasing for 1).
Noting that a > 0t, we have lim h(a) Hence, h(aa) < 0 when a 0. In this way, we have proven that (a( < 0, Le., A (y11.3))) increases as decreases.
G.S. Proof of Proposition 4.2
Restate of Proposition 4.2 The updates induced by a--vdiivegeence for within one gradient descent step are given by:
ply)a ()))))
LogR (a.0) (v) S ())
a
(a)
(b)
p(k)
(k))
+ 9o (v) Σ (k)ss
+ | (a (z)l
(6)
where No.9(y) denotes constant normalization factor independent of y-
33

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a- a-B-Divergence
Proof. The formula for the a-B-divergence is:
D (a.. (p
Σ
p(k)
a g a (k)aat
3
(124)
Using the chain rule, we have:
OD (0 (p q) dy(k)
1l (p (p
Σ
(125)
0f
(k( Ofv
where
0D ()
(p 4)
-
(p()))(())) -
1
(126)
0g(y)
and
04(y) a etv
(127)
05, 0f, Σn etu
Combining Lem. G.2, the Eq. 125 can be expressed as
0
(D 9) 0q(k)
+ @D.(p || q) 0q(v)
Da (l(ll)))=
4)
Σ
Ofy
0g(k) 0f(y)
dq(y) Of,
kzv
(128)
= a
q(v) (p(y) (()))) + q(v)
(aak) -D(K)o)
Now, consider using the gradient descent method to update the loss function with respect to the logits fy. then the distribution the next step y I is given by:
exp(fx+1)
4+ (1)
Σn exp(S
explf; - -1 5;(r
(129)
exp(sh - v) st
exp(-n) ,,+,,
= =
Σ. (kk exp(
Now, substituting the gradient formula of the a- o---ddiveggence, the characterization of d:+1 (0.8) (Ix)) is obtained as:
exp ( [a()) (p(y)* - g(v))) + ((v) (Ex
((4))))))))
444
(a.88 (3) = (v) Σ. (1) exp ( (a()) (p(1) (()))) q(1) (xx ((b)) -()))))((()((]]]]
(130)
Observing that the denominator serves as normalization constant, this can be rewritten as:
4r41 (y) x exp
(n.ft)
((w)s - + a(v) Σ (()))(((())) - - ))
(131)
gr (v)
Taking the logarithm on both sides, we get:
log (a. (v)
77 qc(y)) [ply)= - q(1))))
+ (c(v)) Σ (k)) ( [kc(k) (ook))-jop p(k)
+N!a.9)()) (g),
(1)
a
a
(132)
where N. ((...)(1)) (y) denotes constant normalization factors independent of y. We can further derive that
log
ltt1
(a.9) (a.9)(v) (v)
s ())(
p(y)
(1)
+qc(v) Σ ge (
p(k)
ge
(k:)aa
I
((...9)
(v)
(133)
(v)
a
(a)
(a)
This completes the proof.
34

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via cr- Divver-Dee
G.6. Proof of Theorem D.I
(y) be the distribution obtained after one gradient step, starting from qe using the a-B-divergence. Define A (a,8) as the difference of log mass ratios across two classes y1 and y2, obtained from the a-B-divergence:
Restate of Theorem D.1. Let ge+1 (y)
A ((. (vI. V2) - LogR,
(a) (vi) - LogR (0.3)
(o.) (12))
We have the following (for appropriate positive constants Ç, 8n, 82, and any real numbers on and 02 the range [0, 1] satisfying on o2)::
1. a-B--ivergence transfers probability mass from overestimated classes underestimated classes more aggressively as decreases. yı and y2 are such that ô1 qo(y1) - (e(2)) S p(y)) (where 81 0), and p(y)) 2 p(y2)
holds that A, (a).9) (y1. ( 24 (0
(21,
2. a-B-divergence reduces the probability mass of classes with larger error (p(u) qe(v) more aggressively decreases. If y1 and are such that p(y1) < (0(1)) = (1(22) $ 1 - 62 (where 82 > 0), and p(m)) 2 ply2) + Ç. it holds that A, (01,8) (y1, V2) N A, (a2.9) 92)-
3. The a-B-divergence becomes more (less) preferential in focusing the error on classes with higher student confidence as 3 increases (decreases) when reducing (y) (ogR(a..))))))
Proof. First, we prove Case 1. Note that (3()) = qe (322 = q(x), Based on Eq. 132, we have
A (a.9) (x) p(y1) -(322)
ppls)) 1a-1 dl.
(134)
= =)(())
.p(
Taking the derivative with respect to a, we get
i
A
(a.9)
ng(x)
prine) ra-l In
(135)
Note that p(v1) < and p(12) S 1, we have
a ()
A < O. da
(136)
since In < when E (0, 1]. Therefore, we have A (8) A (o2.9)
when 01 < 02
The proof of Case similar to Case and thus is omitted.
Finally, we prove Case 3. Recall that when reducing
LogR (a (v) we have the following relationship:
LogR(a (a.9) (v) < a
p(y)"
+ (v) Σ (kkss (b)
(2))))
p(k) - (k)))
+ +|N +N
(a.0)
(v)).
k
(b)
where terms (a) and (a1) act as weighting functions. Therefore, selecting larger (smaller) /3 will place more (less) emphasis on errors from classes with higher student confidence qr(y), shown in Fig. 1(c).
Remark. Case and Case 2 show that selecting smaller a leads to a more aggressive reduction of errors across classes by shifting the probability mass from overestimated to underestimated classes. On the other hand, Case shows that increasing emphasizes minimizing errors more classes with higher student confidence.
35

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-
Algorithm 1 alpha beta Divergence Function
Input:
Student distribution qo, Teacher distribution p, and Hyperparameters and / Output: Divergence value ( (pllge)
1: return a. Ex ((()) atpp(k)ats
0 a4za(K)at
Algorithm 2 Generalized distillation framework with divergenee
Input: Dataset D with input-targe target pair {< )n Teacher fr. Student fs. loss weight X, a- a-vv-divereence function ((a..) in Algo. 1, and Hyperparameters and | Output: Trained student model fs
1: for each (In: yo) in D do
2: fo + fr( f t fs((
3: p softmax
4: qo - softmax(f*)
5: (KD + (pllg)) (plle))
6: Update fs towards minimizing (CE(Un;90) + XpxD(P, 90)
7: end for
H. Algorithm Protocol
Algo. 1 and Algo. 2 give the algorithmic protocol of our framework, which is easy implement and applicable to common KD downstream tasks.
I. Additional Experiment Settings
In this section, we provide a more detailed description the experimental protocol.
LL. Natural Language Processing Tasks
1.1.1 DATASETS
Following Gu et al. (2024a); Ko et al. (2024), we select 14K samples from databricks-delly-l (Conover et al., 2023) for training and 500 samples each for validation and testing. After distillation, the models are evaluated on five task- agnostic instruction-following benchmarks: Dolly-evaluation, Self-Instruct, Vicuna-evaluation, Super-Natural Instructions, and Unnatural Instruction. The details for each dataset
follows:
databricks-dolly-15. (Conover et al., 2023): An open-source dataset of instructionn-oolloiing records created by thousands of Databricks employees. It includes several behavioral classes from Ouyang et al. (2022), such as brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.
Self-Instruct (Wang et al., 2023): A framework that enhances a language model's cestuuction-oollwinng by using its own outputs generate extensive instructional data. It contains 52K instructions and 82K input-output pairs for tuning, 252 expert-written tasks for practical use, and 50K public dataset examples for benchmarking.
Vicuna: Utilizes 80 challenging questions for evaluating Vicuna, following (Ko et al., 2024; Gu et al., 2024a).
Super-Natural Instruction (Wang et al., 2022a): benchmark of 1,616 diverse NLP tasks with expert instructions, covering 76 task types. The test set includes 9K samples across 119 tasks.
Unnatural Instruction (Honovich et al., 2023): Al-generated dataset with 240K instructions created with minimal human input, proving AI data can match human data for training language models. The core set has 60K samples.
For experiments on these datasets, we use ROUGE-L (Lin, 2004) as the evaluation metric. ROUGE-L measures the quality of generated text by calculating the Longest Common Subsequence (L.CS) between the generated text y and the reference
36

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--Divrreecce
text ac A higher ROUGE-L score indicates that the generated text is more similar to the reference text. The metric is computed based on the harmonic mean of recall RLcs and precision PLCs, defined as:
LCS(ar, y)
RLcs
Lu
LCS( y) Ly
PLes
ROUGE-L- 2- RiAs Pics
RLCS + PLeS
Here, LCS( y) is the length of the longest common subsequence, Le is the length of the reference text, and Ly is the length of the generated text.
1.1.2. COMPETITORS
Here we give a more detailed summary of the competitors mentioned in the experiments and their SGOs approaches (if exists).
SFT supervised fine-tuning of student model using ground-truth on the Fixed dataset (using predefined input-output pairs).
KD (Hinton, 2015) trains the student distribution to mimic the teacher distribution on the Fixed dataset using FKLD.
SeqKD (Kim & Rush, 2016) maximizes the likelihood of high probability sequences generated by the teacher, and can be viewed as SFT on ceacher-genenann outputs.
MinILLM (Gu al., 2024a) trains the student- genertteee sentences (SGOs) and uses an On-policy gradient method. Their distillation object is to minimize the RKLD between the teacher and student distributions.
GKD (Agarwal et al., 2024) uses the generalized Jensen-Shannon divergence (iiss))) (p (g)) (pll)) + (1 })ge) + (1 - (ge (1 + (1 3))e)), training Mixture of datasets, either reacher-generatee ground-truth. and on-policy student-gent generated sequences.
DISTILLM (Ko et al., 2024) uses Skew KL (W(pllap + -1-aa (1)) Skew RKL (D(gellaq) + (1 -p)) and reports the better performing one. They train mixed dataset consisting of fixed outputs and student- generatee outputs. Additionally, they use Adaptive off off-policy method to determine whether to use student-generatn outputs for training based on validation loss, thereby removing noisy SGOs data.
1.1.3. IMPLEMENTATION DETAILS
Training. For training the teacher and student models, we used four RTX 3090 24GB GPUs. Our experimental setup for training LMs on
primarily follows the experimental setup for Ko et al. (2024). We search for the learning rates in (5c-4, Ic-4, 5e-5}, the batch sizes in {4, 8, 16} within the possible maximum batch size for 3090 24GB GPUs, and train these models for 20 epochs. We fully use the distillation loss for the instruction wwwetiiioooooowinngg dataset and language modeling loss for OpenWebTer. (Gokaslan et al., 2019) corpus. The checkpoints each student are selected by the ROUGE-L scores on the validation set. For all teacher-student configurations, we set 0.2 and 3 0.7. Additionally, the cross-entropy loss ICE was not considered to ensure a fair comparison with previous methods.
To ensure fair comparison, for other competitors, we rerun them (with the necessary hyperparameter tuning) and select the best-performing checkpoint on the validation set.
Evaluation. For evaluating the teacher and student models, we applied a single RTX 3090 24GB GPU. Following (Ko et al., 2024; Gu et al., 2024a), We adopt a prompt template as shown in Fig. 7. We sample the responses from each model using a temperature of 1. 1.0, max-length limit of 512, and five random seeds (i.e., { (10, 20, 30, 40, 503).
37

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3-DDvvereecce
Below is an instruction that describes task. Write a response that appropriately completes the request.
### Instruction: (instruction)
### Input:
(input)
### Response:
Figure 7. The prompt template for training and evaluation of instruction inslluction-ollooiing task experiments from (Ko et al., 2024; Gu et al., 2024a).
1.2. Vision Tasks
1.2.1. DATASETS
In this section, we provide detailed descriptions of the image datasets used.
It is generic image dataset consisting of 60,000 32x32 color images across 100 classes, with 600 images per class. It is further split into 50,000 training images and 10,000 test images.
CIFAR-100:
ImageNet (Deng et al., 2009): It is widely recognized object classification dataset containing approximately 1.28 million training images and 50,000 test images across 1,000 object classes. The images are sourced from the web and organized using the WordNet hierarchy, making it standard benchmark for evaluating object recognition models.
Caltech101 (Fei-Fei et al., 2004): It is an object classification dataset with 101 classes and a background class, containing approximately 7,650 training images and 3,300 test images. The images vary significantly in scale, orientation, and lighting conditions.
OxfordPets (Parkhi et al., 2012): It is a fine- ane-ggrined pet classification dataset with 37 pet breed classes, featuring nearly equal numbers of training (3,680) and test (3,669) images. It also includes pixel-level segmentation masks.
StanfordCars (Krause et al., 2013): It is fine-grained car model recognition dataset with 196 classes, based on make, model, and year. It contains 8,144 training images and 8,041 test images, capturing diverse vehicle angles and environments.
Flowers 102 (Nilsback & Zisserman, 2008): It is dataset of 102 flower species for fine-grained classification tasks. It includes 6, 149 training images and 1,020 test images, posing challenges in distinguishing visually similar flower classes.
Food101 (Bossard et al., 2014): is fine-grained food classification dataset with 101 dish classes, comprising 75,750 training images and 25,250 test images. It poses challenges in recognizing overlapping ingredients and presentation styles.
FGVCAircraft (Maji al., 2013): It is fine-grained aireraft classification dataset with 100 classes, distinguishing between models and manufacturers. It contains 6,667 training images and 3,333 test images.
SUN397 (Xiao et al., 2010): is comprehensive scene recognition dataset with 397 classes, including natural landscapes, indoor spaces, and urban environments. It contains approximately 50,000 training images and 50,000 test images.
UCF101 (Soomro, 2012): It is video dataset for action recognition, featuring 101 action classes ranging from sports to daily activities. It contains approximately 9,500 training clips and 3,700 test clips, collected from YouTube.
DTD (Soomro, 2012): is a texture classification dataset with 47 texture classes described using human- iamaaninterprectbble attributes. It contains 3,760 training images and 1,880 test images.
38

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--Divergence
EuroSAT (Helber et al., 2019): It is satellite image dataset for land-use and land-cover classification, with 10 classes such as agricultural areas, forests, and urban regions. It contains 21,600) training images and 5,400 test images.
For experiments these datasets, we follow the popular setup in classification tasks to accuracy the evaluation metric.
1.2.2. COMPETITORS
In this section, we provide a more in-depth overview of the competitors discussed in the vision experiments.
First, we introduce the distillation-baser methods:
KD (Hinton, 2015) directly minimizes the FKLD between the student and teacher distributions transfer knowledge.
DKD (Zhao et al., 2022) uses FKLD for distillation, where the knowledge of the teacher's distribution is decoupled into target class and non-target class knowledge for separate learning.
LSD (Sun al., 2024) uses FKLD for distillation in their experiments, where they first normalize the logit vector before obtaining the model output distribution.
TTM (Zheng & Yang. 2024) uses FKLD for distillation, and they also introduce Rényi entropy as a regularization to make the student distribution smoother.
Next, we describe the SFT-based methods:
CLIP supervised fine-tuning using ground-truth on standard datasets.
CoCoOp (Zhou et al., 2022) enhances new class performance by transforming the unified context into an instance- adaptive context, where each sample is assigned specific prompt that focuses its unique features attributes.
MaPLe (Khattak et al., 2023a) improves vision-language alignment by simultaneously adapting both the text and image encoders in CLIP using hierarchical prompts.
PromptSRC (Khattak et al., 2023b) ensures better performance on both base and new classes by minimizing the task cross-entropy loss and the FKLD between the output distribution of the model and the pre-trained model.
1.2.3. IMPLEMENTATION DETAILS
We conduct all vision experiments on a single RTX 3090 GPU. The detailed experimental setups are as follows.
setup. In this experimental setup, we consider model architectures including VGG (Simonyan, 2014), ResNet (He et al., 2016), and WideResNet (Zagoruyko, 2016). Following (Zheng & Yang, 2024; Sun al., 2024; Zhao et al., 2022), we train the student models on all class samples. We also consider a standard training data augmentation scheme including padding 4 pixels prior random cropping and horizontal flipping. We set the batch size as 64 and the initial learning rate as 0.05. We train the model for 240 epochs, in which the learning rate is decayed by 10 every 30 epochs after 150 epochs. We use stochastic gradient descent (SGD) as the optimizer with weight decay 5e-4 and momentum 0.9.
Standard Training-Evaluatio
For evaluation, we report the average accuracy across all classes on the test set. We list the hyperparameters a and 3 used across the above experiments in Tab, 4. In our ABKD, the weight À of ÉKD is set to the default value 32. For those re-implemented methods, we only adjust and 8 and follow the other hyperparameters reported in their original papers.
Base-to-New setup. In this experimental setup, we use the VIT-L/14 CLIP model as the teacher and the VIT-B/16 CLIP model as the student. We adopt the recently popular prompt tuning setup for CLIP, as it performs sufficiently well across many tasks, despite freezing most of the model parameters and training only a subset of the learnable prompt tokens. We split the training and testing datasets into base and new classes same as previous work (Khattak et al., 2023a; Kim et al., 2024). Tab. 5 provides the details of the number of images used for training on the base-to-new setup.
39

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3-Divergence
Table 4. Hyperparameters for different architecture distillations CIFAR-100.
Teacher Student
WRN-10-2
WRN-10-2
resnet56
resnet110
resnet110
respet32xe
veg13 vgg8
resnet110
WRN- 16-2
WRN-40-1
resnet20
resnet20
resnet32
resnetBx4
resmet44
ABKD
0.
09, - 0.2
0.8,2
0.8,8
0.3
0.7.6
0.9, 0.2
0.6,3 0.3
ABDKD
0.8.8
0.4
1 1 0.2
a- 100,3
0.2
a- 0.8,8
0.3
0.9,5,5 -0.
0,8,3 3 0.3
a- 0.7,2
0.4
a- 0.8,3 0.3
o=
0.9.3
0 04
a= 09,8
o= 12,,
0.1
0.009,6
122,8
-02
a= 1.0, 0.2
100,8
ABTTM
0.8,9 0.3
1 0.1
0.7,2
0.5
o= 0.9,0
00,,00000
0 0.3
0.7, = 0.5
0.8,3 0.2
Table Number of images used for distillation and testing per-dataset. To ensure fair comparison, follow the data split as prior
arts (Kim et al., 2024).
Dataset
ImageNet
Caltech101
OxfordPets
StanfordCars
Flowers102
Food101
FGVCAireraft
SUN397
DTD
EuroSAT
UCF101
Train
1281,167
4,128
2,944
6,509
4,093
50,500
3,334
15,880
2,820
13,500
7,639
Test Base
25,000
1,549
1,881
4,002
1,053
15,300
1,666
9,950
864
4,200
1,934
Test New
25,000
916
1,788
4,039
1,410
15,000
1,667
9,900
828
3,900
1.849
The teacher is pre-trained using the PromptSRC (Khattak et al., 2023b) method, following which it is fine-tuned on the base classes using ground truth supervision. Following (Kim et al., 2024), all distillation-basee methods use the unlabeled training to train students for a fair comparison, and we search for in {100, 200,300, 500, 1000, 2000, 3000]. We set the prompt depth to and the vision and language prompt lengths to 4. We SGD as the optimizer. All student models are trained for 20 epochs with batch size of 8 and learning rate of 0.005. We follow the data augmentation scheme as in Khattak et al. (2023b), i.e., random resized cropping and random flipping. The text prompts of the first layer are initialized with the word embeddings of "a photo of a {classname)".
For evaluation, we report the model's accuracy on both the base classes and the new classes separately. Additionally, we report the Harmonic Mean (HM) of the two accuracies, defined as:
2 x BaseAcc x NewAce
HM
(137)
BaseAcc NewAcc
We list the hyperparameters a and used across different datasets in Tab. 7. In addition, since LSD, DKD, and TTM did report performance under the base-to-new setting, we reran their source code and report the best results (with necessary hyperparameter tuning).
Table 6. Hyperparameters for different image datasets
Dataset
ImageNet
Caltech101
OxfordPets
StanfordCars
Flowers
Food101
FGVCAireraft
SUN397
DTD
EurOAAT
UCF101
0.5
0.8
0.8
0.6
0.9
0.5
0.6
0.8
1.0
0.6
0.8
8
0.5
0.2
0.4
0.4
0.1
0,5
0.5
0.2
0.2
0.5
0.2
J. Additional Experiment Analysis
In this section, we present additional experimental results on language and vision tasks.
J.L. Natural Language Processing Tasks
J. EFFECTS OF SGOS
Prior arts (Ko et al., 2024; Agarwal al., 2024) highlight that existing KD methods suffer from distribution mismatch between the output sequences seen during training and those generated by the student during inference in auto-regressive language models. To address this, these works incorporate (i.e., token-level predict distribution for these sentences) during training, leading to significant improvements. To assess the applicability our framework, we evaluate the performance after training with different SGO strategies, as shown Tab. 8. The results indicate that by integrating these promising techniques, our framework achieves further improvements across most datasets compared to training with fixed data.
outputs (SGOs) along with teacher feedback
40

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a--3-
Table 7. Hyperparameters for different instruction datasets
Dataset
Dolly Eval
Self-Instruct
Vicuna Eval
Super-Natural
Unnatural
0
0.2
0.2
0.2
0.2
0.2 0.7
8
0.7
0.7
0.7
0.7
Table 8. Effects of our framework using different SGOs strategies (ie. On-policy, Mixed, and Adaptive Off-policy). Fixed denotes that our framework uses only the original dataset for training without augmentation. Following (Ko et al., 2024), in the mixed strategy, we apply the on-policy optimization method with a probability of 0.5. Otherwise, sample from the fixed dataset.
Method
Dolly Eval
Self-Instruct
Vicuna Eval
Super-Natural
Unnatural
Prior SOTA result
25.32 (0.14)
12.49 (0.56)
17.30 (0.41)
23.76 (0.38)
25.79 (0.08)
Fixed + Ours
25.65 (0.24)
13.47 (0.42)
16.06 (0.25)
26.47 (0.31)
29.32 (0.08)
On-policy (Gu et al., 2024a) + Ours
25.96 (0.42)
13.44 (0.37)
17.32 (0.38)
26,86 (0.26)
29.57 (0.13)
Mixed (Agarwal et al., 2024) + Ours
26.49 (0.23)
14.62 (0.27)
17.14 (0.26)
27.54 (0.44)
30,98 (0.09)
Adaptive Off-policy (Ko al., 2024) + Ours
26.58 (0.18)
14.25 (0.25)
17.79 (0.35)
27.79 (0.26)
31.13 (0.12)
J.1.2. DISTILLING FROM STRONGER TEACHER
Recent research (Cho& & Hariharan, 2019; Huang et al., 2022) shows that as the size of the teacher model increases, the distillation performance does not always improve for the student models and may even degrade due to the capacity gap between them. It is not clear how our framework performs when scaling up the teacher models' sizes. To this end, we report the performance of our method using teacher models of varying sizes while keeping the student model size fixed, as shown in Tab. 9. From the results, we have the following observations: 1) Our a----divergence consistently outperforms FKLD and RKLD across different teacher models; 2) FKLD and RKLD fail to ensure that the student model consistently benefits from the rich supervision provided by larger teacher models and thus leads suboptimal performance. 3) In contrast, the a- a-3-divergeence can maintain the student model's performance nearly positively correlated with teacher model size by smoothly interpolating between FKLD and RKLD.
Table 9, Performance of GPT-2 on five task agostic instruction Instruction-folloiinn datasets with different teacher model sizes.
Method
Dolly Eval Self-Instruet
Vicuma Eval
Super-Natural
Unnatural
GPT-2 Medium (0.3B)
GPT-2 (0.IB)
FKLD
23.68 (0.29)
10.14 (0.53)
15.44 (0.48)
18.54 (0.30)
20.44 (0.20)
RKLD
24.66 (0.20)
11.73 (0.31)
15.27 (0.41)
22.65 (0.25)
25.27 (0.24)
a-S-divergence (Ours)
25.47 (0.25)
13.13 (0.46)
15.84 (0.21)
26.29 (0.13)
28.31 (0.14)
GPT-2 Large (0.8B)
GPT-2 (O.IB)
FKLD
20.01 (0.23)
9.89 (0.59)
14.98 (0.35)
19.00 (0.26)
18.72 (0.13)
RKIDD
25.27 (0.28)
11.77 (0.24)
14.78 (0.26)
23.61 (0.36)
26.41 (0.13)
a-S-divergence (Ours)
25.72 (0.52)
13.08 (0.45)
15.80 (0.62)
26.44 (0.32)
29.25 (0.10)
GPT-2 XL (1.58) + GPT-2 (B.1B)
FKLD
23.80 (0.37)
10.01 (0.75)
15.25 (0.65)
17.69 (0.26)
18.99 (0.05)
RKLD
24.77 (0.37)
12.02 (0.48)
15.06 (0.28)
23.27 (0.29)
26.01 (0.11)
a-U-divergence (Ours)
25.65 (0.24)
13.47 (0.42)
16.06 (0.25)
26.47 (0.31)
29.32 (0.08)
J.1.3. COMPARISON WITH MORE BASELINES
To further demonstrate the effectiveness of the proposed method, we additionally consider several KD baselines that were not included in the main text, namely (1) AlphaNet (Wang et al., 2021), (2) BDKD (Amara al., 2022), (3) AKL (Wu et al., 2024), and (4) Jensen's KL (Binici et al., 2022).
The results are shown in Tab. 10. ABKD outperforms all other baselines across different benchmarks, with improvements ranging from 0.42 to 1.76. In addition, we visualize the performance dynamics of different methods throughout the training process, as shown in Fig. 8. The a-3 divergence consistently achieves the highest performance throughout training, especially in the early stages, demonstrating its faster convergence ability.

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via 1e-3--Divergence
28.0
Val ROUGE-L
26.0
24.0
22.0
PKLD
Jensen's KL
RK.D
WED
AlphaNet
18.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
epoch
Figure Performance on the validation set when distilling GPT-2 XL (1.5B) to GPT-2 (0.1B).
Finally, it is worth noting that although AlphaNet achieves performance comparable ours, it requires tuning three hyperparameters (while ours only has two), which significantly increases the search overhead. In addition, key advantage of our approach is that its hyperparameters have clear physical interpretations. This allows one to leverage inductive bias and follow principled guidelines to more efficiently search for suitable hyperparameters for the target task (App. D).
Table 10. ROUGE-L scores (f) of different loss functions on five task- aagnosicciicc instruction cestruction-folowiing datasets when distilling GPT-2 XL (1.5B) into GPT-2 (B.1B). We report the average and standard deviation ROUGE-L scores across five random seeds [10, 20, 30, 40,
50J.
Loss Function
Dolly
Self-Instruct
Vicuna
Super-Natural
Unnatural
FKLD
23.80 (0.37)
10.01 (0.75)
15.25 (0.65)
17.69 (0.26)
18.99 (0.05)
RKLD
24.77 (0.37)
12.02 (0.48)
15.06 (0.28)
23.27 (0.29)
26.01 (0.11)
WSD
23.33 (0.52)
10.52 (0.47)
14.83 (0.61)
19.67 (0.13)
21.21 (0.21)
BDKD
23.94 (0.24)
11.83 (0.39)
15.21 (0.23)
19.56 (0.23)
21.66 (0.23)
Lensen-Shannor divergence
23.79 (0.24)
11.52 (0.18)
15.35 (0.80)
21.36 (0.17)
21.97 (0.10)
AKL
23.83 (0.59)
10.87 (0.42)
15.63 (0.66)
20.07 (0.32)
21.97 (0.13)
AlphaNet
25.13 (0.27)
12.46 (0.46)
15.64 (0.40)
25.27 (0.20)
27.56 (0.15)
SKL
25.01 (0.23)
12.47 (0.29)
15.98 (0.84)
25.56 (0.31)
27.51 (0.07)
SRKL
25.75 (0.39)
11.58 (0.49)
15.56 (0.17)
26.13 (0.25)
27.37 (0.18)
a- a-divergence
25.15 (0.41)
12.92 (0.22)
15.60 (0.27)
24.83 (0.21)
27.81 (0.10)
8- d-divergence
24.12 (0.38)
11.18 (0.27)
14.95 (0.33)
20.98 (0.23)
23.15 (0.14)
a- --diierreeece (Ours)
25.65 (0.24)
13.47 (0.42)
16.06 (0.25)
26.47 (0.31)
29.32 (0.08)
J. 1.4. LLAMA FAMILY DISTILLATION
The following analysis aims to investigate whether the proposed method remains effective in distillation experiments involving larger-scale models. To this end, we conducted distillation from OpenLLaMA2-7B (Touvron et al., 2023b) 3B and compared approach with various KD baselines.
The results are presented in Tab. 11. ABKD outperforms others by 0.65-3.26 ROUGE-L scores, especially excelling in Dolly and Unnatural.
J. QUALITATIVE EVALUATION
In this section, we present several case studies to illustrate the effectiveness of ABKD, As shown in Tab. 17, ABKD better generating more accurate responses according to the predefined requirements specified in the instructions.
42

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a----Divergene
Table ROUGE-L scores (1) on task-agnostic instruction instluctionfollooinngg datasets when distilling OpenL into Ope2L
3B. Experiments are conducted on eight RTX 3090 24GB GPUs. indicates that SGOs are used.
Method
Dolly
Self-Instruct
Vicuna
Super-Natural
Unnatural
SFT
24.54 (0.51)
16.80 (0.64)
16.15 (0.15)
29.29 (0.13)
27.43 (0.21)
FKLD
25.23 (0.44)
18.90 (1.20)
16.67 (0.35)
31.68 (0.22)
29.36 (0.13)
RKLD
27.74 (0.45)
20.61 (0.80)
18.83 (0.40)
35.31 (0.24)
33.86 (0.16)
Jensen's KL
26.28 (0.43)
18.84 (0.66)
17.81 (0.38)
30.92 (0.12)
29.79 (0.17)
BDKD
26.78 (0.53)
18.94 (0.68)
17.81 (0.52)
32.15 (0.34)
30.89 (0.24)
AKL
26.38 (0.41)
17.69 (0.46)
16.72 (0.48)
33.02 (0.16)
31.29 (0.08)
DISTILLM"
28.24 (0.48)
21.00 (0.72)
19.12 (0.53)
37.06 (0.35)
35,05 (0.13)
AlphaNet
28.11 (0.29)
21.30 (0.63)
18.70 (0.23)
37.86 (0.44)
35.40 (0.17)
Ours (ABKD)
30.25 (0.37)
22.39 (0.62)
20.83 (0.42)
38.51 (0.32)
38.66 (0.10)
Table 12. The distillation results of Qwen2.5-Math English mathematical benchmarks. Models evaluated with chain-of-thought prompting.
BENCHMARK
GaoKao 2023 En
Olympiad Bench
College Math
GSM8K
MATH
Avg.
MODEL
Qwen2.5 1-..55--Mt--BBBBaaa IB-Instuutt (Teacher)
95.5
82.8
66.8
38.5
37.7
64.3
Qwen2.5-1.5B-Innst Instrruct (Student)
73.3
54.9
45.9
18.9
30.3
44.7
SeqKD
75.8
57.3
47.3
17.7
31.3
45.9
KD
75.9
58.1
45.5
21.1
31.3
46.3
ABKD
77.4
58.6
48.5
20.4
32.0
47.4
J.2. Vision Tasks
1. DISTILLING FROM STRONGER TEACHER
To evaluate the potential our framework to benefit from larger teacher, we examine the distillation effect when different- sized teacher models are used to distill student model of the same size, as shown in Tab. 13. Encouragingly, our framework consistently outperforms FKLD and RKLD, with performance gains remaining stable as the teacher model size increases.
Table 13. Performance resnet20 on CIFAR-100 with different teacher model sizes. a fair comparison, we set a and , = 0.3
across all leacher-student configurations.
Accuracy (%)
Student
Teacher
Student
Teacher
FKLD
RKLD
Q- o-B-divergence (Ours)
resnet32
71.93
71.03 (0.23)
70.91 (0.29)
71.46 (0.15)
resnet44
72.25
71.51 (0.11)
71.29 (0.16)
71.76 (0.25)
resnet20
69.06
resnet56
72.34
70.66 (0.24)
71.43 (0.16)
71.79 (0.16)
resnetl10
74.31
70.67 (0.27)
71.41 (0.23)
71.72 (0.18)
J.2.2.
DISTILLATION
Although the analysis in the main text has demonstrated the effectiveness of the proposed method for distillation within the same architecture, it remains unclear how much improvement our method can achieve when the teacher and student have different architectures. To this end, we performed distillation from ResNet50 to VGG8, The results in Tab. 14 show that our method outperforms previous approaches by margin of 0.15 to 0.89.
43

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via a-3--Divrrgene
Table 14. Accuracy (%) comparison of different distillation methods from ResNet50 VGG8 CIFAR-100.
Method
Accuracy (%)
KD
73.81
ABKD (Ours)
74.62 (0.81)
DKD
74.37
ABDKD (Ours)
75.26 (0.89)
LSD
74.52
ABLSD (Ours)
74.77 (0.25)
TTM
74.87
ABTTM (Ours)
75.02 15)
J.2.3. HOW DOES ABKD PERFORM WITH ALPHA/BETA OUTSIDE [0,1]?
Another interesting question how ABKD performs when a or / fall outside the range [0, (e.g., a 1.5, = -0..).
To address this concern, tested the settings with and 3 < 0 when distilling ResNet56 ResNet20, shown in Tab. 15.
Table Accuracy (%)
ABKD with a > and 3 < when distilling ResNet56 ResNet20.
a
-0.1
-0.3
-0.5
1.2
70.81
71.10
70.35
1.4
71.29
71.24
70.92
1.6
70.55
70.53
70.34
The results indicate that excessively large values a weaken the condeess-conceenraation effect, while overly small values of diminish the ceffidenco-concentaatiio effect. Both cases can lead to degraded distillation performance. This observation aligns with our objective: we aim to balance FKLD and RKLD, which correspond to the extreme cases 1, a 0, and
1. Therefore, a natural approach is search for parameters within the range [0, 1].
44

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via c- a-3-Divergence
Table 16. Comparison with existing SOTA methods on base-to-new generalization. Teacher: ViT-L/14 CLIP. Student: ViT-B/16 CLIP.
ViT-B/16
Base
Novel
HM
ViT-B/16
Basc
Novel
HM
ViT-B/16
Base
Novel
HM
Teacher
87.85
81.45
84.39
Teacher
83.24
76.83
79.91
Teacher
98.71
98.03
98.37
CLIP
69.34
74.22
71.70
CLIP
72.43
68.14
70.22
CLIP
96.84
94.00
95.40
CoCoOp
80.47
71.69
75.83
CoCoOp
76.66
70.54
73.47
CoCoOp
97.96
93.81
95.84
MaPI.e
82.28
75.14
78.55
MaPL.e
77.00
74.05
75.49
MaPLe
97.74
94.36
96.02
PromptSRC
84.26
76.10
79.97
PromptSRC
77.63
74.97
76.26
PromptSRC
98.10
94.03
96.02
KD
86.96
80.73
83.63
KD
80.83
74.66
77.62
KD
98.91
96.65
97.77
DKD
87.02
81.02
83.79
DKD
80.98
74.85
77.79
DKD
99.12
96.52
97.80
LSD
86.31
79.99
82.89
LSD
80.85
74.62
77.61
LSD
99.05
96.24
97.62
Ours
87.27
81.41
84.17
Ours
81.23
75.02
78.00
Ours
99.46
96.93
98.18
(a) Average over 11 datasets
(b) ImageNet
(c) Caltech101
ViT-B/16
Base
Novel
HM
ViT-B/16
Base
Novel
HM
ViT-B/16
Base
Novel
HM
Teacher
96.86
98.82
97.83
Teacher
84.53
84.25
84.39
Teacher
99.05
82.60
90.08
CLIP
91.17
97.26
94.12
CLIP
63.37
74.89
68.65
CLIP
72.08
77.80
74.83
CoCoOp
95.23
97.69
96.43
CoCoOp
70.49
73.59
72.01
CoCoOp
94.47
71.75
81.01
MaPL.e
95.43
97.96
96.68
MaPLe
72.94
74.00
73.47
MaPLe
95.92
72.64
82.56
PromptSRC
95.33
97.30
96.30
PromptSRC
78.27
74.97
76,56
PrompiSRC
98.02
76.50
85,92
KD
96.30
98.01
97.15
KD
82.80
83.37
83.13
KD
99.42
82.62
90.24
DKD
96.36
98.52
97.43
DKD
82.23
84,20
83.21
DKD
99.15
82.64
90.15
LSD
95.96
98.32
97.13
LSD
78.29
79.48
78.88
LSD
98.86
81.84
89.55
Ours
96.49
98.55
97.51
Ours
83.43
84.01
83.72
Ours
99.24
83.47
90.67
(d) OxfordPets
(e) StanfordCars
(f)
Flowers102
ViT-B/16
Novel
HM
ViT-B/16
Novel
HIM
ViT-B/16
Novel
HM
Teacher
94.56
95.15
94.85
Teacher
$4.44
43.07
48.09
Teacher
84.97
81.09
82.98
CLIP
90.10
91.22
90.66
CLIP
27.19
36.29
31.09
CLIP
69.36
75.35
72.23
CoCoOp
90.70
91.29
90.99
CoCoOp
33.41
23.71
27.74
CoCoOp
79.74
76.86
78.27
MaPLe
90.91
91.25
91.08
MaPLe
37.44
35.41
36.50
MaPLe
82.88
78.70
80.75
PromptSRC
90.67
91.53
91.10
PromptSRC
42.73
37.87
40.15
PromptSRC
82.67
78.47
80.52
92.43
93.68
93.05
49.12
41.81
45.17
83.69
81,54
82.60
DKD
92.35
93.72
93.03
DKD
48.92
42,43
45.44
DKD
83.87
81.32
82,58
LSD
92.07
93.07
92.57
LSD
47.76
39.84
43.44
LSD
83.34
80.62
81.96
Ours
92.46
93.84
93.14
Ours
49.06
43.05
45.86
Ours
83.88
81.75
82.80
gg) Food101
(h) FGVCAireraft
SUN397
ViT-B/16
Base
Novel
HM
VIT-B/16
Base
Novel
HM
VIT-B/16
Base
Novel
HM
Teacher
85.76
70.65
77.48
Teacher
94.79
83.15
88,59
Teacher
89.50
82.26
85.73
CLIP
53.24
59,90
56.37
CLIP
56.48
64.05
60.03
CLIP
70.53
77.50
73.85
CoCoOp
77.01
56.00
64.85
CoCoOp
90.89
71.64
80.11
CoCoOp
82,33
74.95
78.47
MaPLe
80.36
59.18
68.16
MaPLe
93.32
71.26
80.94
MaPLe
83.08
78.60
80.76
PromptSRC
83.37
62.97
71.75
PromptSRC
92.90
73.90
82.32
PromptSRC
87.10
78.01
82.30
KD
85.84
71.37
77.94
KD
97.54
82.08
89.14
KD
89.71
82,27
86.10
DKD
86.83
71.91
78.67
DKD
97.14
83.72
89.93
DKD
90.28
81.40
85.61
LSD
86.35
70.97
77.91
LSD
97.31
83.31
89.77
LSD
89.62
81.53
85.38
Ours
86.52
72.65
78.98
Ours
97.75
83.92
90.30
Ours
90.44
82.31
86.73
(i) DTD
(k) EuroSAT
(l) UCF101
45

---

ABKD: Pursuing Proper Allocation of the Probability Mass in Knowledge Distillation via
Table 17. cnprucctionfollowinng cases from the Unnatural (Honovich et al., 2023) dataset. ABKD is better at generating responses that meet the requirements of the given instructions.
Case #1
In this lask, you's 'e given a short description an event. Your job to order the steps involved in the evens from first to last, Note that there may be multiple correct answers for each event.
Instruction
Input
Description: A man walks into a bar and orders a drink. He pays for his drink and leaves the bar:
Ground-truth
The man walks the har.) in He orders drink. \n He pays
drink. (n He leaves the har. A man walks into bar, asks a burtender who the har and drinks.
SFT
A walks bar, pays for his drink, and leaves the bar with short deseription of the event. inn woman walks into coffee shop, opens the door, and asks the bartender for him to order them. \oa 'an A man asks the barteader for their order and alcohol. The woman hands the coffeee and begins making her coffee. At the end, men exit, and woman returns bar.
KD
SogKD
A man walks into bar with drink and leaves the bur: he pays for his drink and enters the bar.
MiniL.IM
The event was man walking into har and ondering a drink.
GKD
A man walks
bar. paid for drink and leaves The
DISTILLM
A man walks into bar, orders drink and leaves the bar
Ours (ABKD)
A man walks into bar orders drink. He pays for his drink leaves the bur.
Case #
Instruction
You need output the text of the passage after capitulizing all proper nouns.
Input
Rami went to new york city on business. While be was there, be met his good friend ali, who showed him around town.
Ground
Rumi went to York City on business. While was there, met his good friend Ali, who showed him around town.
SFT
Rami went to new yORK city on business. While she was there, be met His good friend ali, who shows him around town.
KD
Rami went to new yORK city on business. While was there, met good friend ali, who taught him bow to go city.
SeeKD
Rumi went to new yorks city on business. While being there, he met his had friend ali, who showed you around town.
MiniLLM
Rami went to new yorks city on business. While being there, he mot friead ali, who showed he around town.
GKD
Rami new yORK city business, While
his bad friend Ali, who showed him around Town.
DISTILLM
York City. Holds; ali, who showed him that you are good friead.
Ours (ABKD)
Rani went to New York Citys on business. While was there, be met her good friend Ali, who shown him around town.
Case #3
You will be given a text consisting of multiple sentences. The task is to find the number of questions present the text, and then print them out sepurately. A question defined as sentence that ends with question mark ("T).
Instruction
Input
Text: was on October 3rd, 1984, you know when you were bom? What month are in currently?
Ground truth
Do you know when you were born? \/n What month are we in currendly?
SFT
was bom October 3. 1984, What
in currently?
KD
I was hom on October 3nd, 1984 Do you actually know? I'm not sure because I don't 't know when I was homn.
SeyKD
You in currently.
MiniLLM
The text of the question is: I was bom on September 3rd, 1984,
GKD
am born on October 3rd 1984, Do you know what monch are we in currenily, year?
The task find answer question that begins with question mark (""". The task find answer question that begins with question mark (!"). The task is to find the answer to a question that begins with question mark ('!).
DISTILLM
Ours (ABKD)
was bor on October 3nd, 1984, Do you remember when you were born?
46


Reviews:
[{"review_id": "OWiDDrHb2F", "reviewer": "Reviewer_71pK", "summary": "The paper discusses the main challenges in knowledge distillation, which lies under the proper balance between two modes (1) hardness concentration and (2)confidence concentration. They provided a smoother transition between the reverse and forward KL divergences via the integration of alpha-beta divergence They performed analyses on both language and vision tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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**):\n>Q1: Essential References Not Discussed\n\n**A1:** We will include the discussion of these baselines (the experiments are in Q5) in the new version:\n\n- 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$.\n\n- BDKD adjusts the weights of FKLD and RKLD during training based on the entropy difference between $p$ and $q$.\n\n- 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$.\n\nThe 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. \n\nOverall, 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.\n>Q2: More Details on hyperparams α and β (Weakness #1)\n\n**A2:** Thank you for your insightful question! We agree on the importance of clarifying hyperparam selection.\n\nFirst, 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).\n\nSecond, 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.  \n\nOverall, 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.\n>Q3: Training Cost and Convergence Behavior (Weakness #2)\n\n **A3**: Thank you for your insightful question!\n\n **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.\n\n **Convergence**: Fig.1 in the link shows that our method outperforms others at all training stages (**especially in the early stages**), showing **faster convergence**.\n>Q4: Novelty issue (Weakness #3)\n\n**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:\n\n- FKLD has overly weak hardness- and confidence-concentration effects. It fails to focus on the target class, leading to wrong predictions.\n- RKLD has overly strong hardness- and confidence-concentration effects. It overemphasizes the target class while ignoring the broader knowledge from non-target classes.\n- The weighted sum of FKLD and RKLD tends to overemphasize the minimal values in both teacher and student distributions.\n- α-divergence imposes an unnecessary sum-to-one constraint on hardness- and confidence-concentration effects, potentially limiting model performance.\n- α-β-divergence offers a flexible way to adjust hardness- and confidence-concentration separately, enabling smooth interpolation between FKLD and RKLD.\n\nThese insights are rarely explored in existing literature.\n\n> Q5: Performance Improvement & More Baselines (Weakness #4)\n\n**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. \n||Dolly|Self-Instruct|Vicuna|Super-Natural|Unnatural|\n|-|-|-|-|-|-|\n|Prior SOTA|25.13 (0.27)|12.46 (0.46)|15.64 (0.40)|25.27 (0.20)|27.56 (0.15)|\n|Ours|**25.65** (0.24)|**13.47** (0.42)|**16.06** (0.25)|**26.47** (0.31)|**29.32** (0.08)|  \n\nTab.2 (see the link for the full version) shows our method outperforms others by 0.42-1.76.\n\nTo further validate our method, we made our best effort in the past few days to study **OpenLLaMA-8B→3B** distillation. \n||Dolly|Self-Instruct|Vicuna|Super-Natural|Unnatural|\n|-|-|-|-|-|-|\n|Prior SOTA|28.24 (0.48)|21.30 (0.63)|19.12 (0.53)\t|37.86 (0.44)|35.40 (0.17)|\n|Ours|**30.25** (0.37)|**22.39** (0.62)|**20.83** (0.42)|**38.51** (0.32)|**38.66** (0.10)| \n\nTab.3 (see the link for the full version) shows it outperforms others by 0.65-3.26, especially excelling in Dolly and Unnatural.", "claims_and_evidence": "- The claim of using alpha-beta divergence for better distillation seems to ambiguous and not empirically well supported in the paper. \nIndeed, introducing both alpha and beta as hyper-parameters adds more complexity to the already challenging problem of distillation. \nDifferent alpha and beta can lead to very different distillation behavior. This could eventually impact the stability and the convergence of the student network while performing distillation.", "methods_and_evaluation_criteria": "The paper evaluates their technique on both language and vision tasks, which I think is good and makes their work more comprehensive.", "theoretical_claims": "Yes. I checked the theoretical proofs. I have not noticed any issues.", "experimental_designs_or_analyses": "The experimental design and the analyses is somewhat comprehensive. The authors tried to empirically compare their proposed technique alpha-beta KD to other KD techniques. However, I would like to address certain comments with respect to the experiments:\n- There are some missing baselines. Before balancing the divergences, previous literature have proposed to implement Jensen's KL divergence in which they symmetrically use both the RKLD and FKLD. Adding this as a baseline in their experiments could contrast the performance of the distillation. \n- With respect to analyses, as heavy user of distillation, I would like to see the computational and complexity trade-off that the alpha-beta KD offers. From the results of the paper, the reported improvements over baselines like vanilla KD and other distillation methods appear to be marginal or even lower than that of alternative KD methods. This really raises concerns about the added complexity and actual performance gains. It would have been better if the authors could add other analyses and justification of why this method would be preferable despite the noticeable limited improvements.", "relation_to_broader_scientific_literature": "The key contribution of the paper closely resembles existing techniques that have already been explored and discussed in earlier literature on knowledge distillation. The proposed approach does not introduce a fundamentally novel concept but rather builds upon previously established methods.", "essential_references_not_discussed": "Some key references to consider are earlier works that explored the use of Jensen-KL (symmetric KL) for knowledge distillation. [1]\n\nAdditionally, there are prior studies on adaptive KL divergence, which, similar to the idea presented in this paper, aim to balance forward KL (FKLD) and reverse KL (RKLD). A proper discussion and citation of these works would provide better context and highlight the connections between this study and existing research. [2] \n\n\n[1] Binici, Kuluhan, et al. \"Preventing catastrophic forgetting and distribution mismatch in knowledge distillation via synthetic data.\" Proceedings of the IEEE/CVF winter conference on applications of computer vision. 2022.\n\n[2] Amara, Ibtihel, et al. \"Bd-kd: balancing the divergences for online knowledge distillation.\" arXiv preprint arXiv:2212.12965 (2022).", "other_strengths_and_weaknesses": "I would like to summarize what I have mentioned above:\nStrengths: \n- The paper is well written and easy to follow.\n- The paper evaluates on various language and vision tasks. \n\nWeaknesses:\n- It would have been better if the authors could justify more about the hyper-parameters of alpha and beta. While the authors specify the values they used in their experiments, it remains unclear how extensively they tuned these parameters to achieve optimal performance.\n- Another key aspect missing is the discussion on training time. Based on my expertise, training with alpha divergence tends to have slower convergence compared to traditional KD losses. Aside from the added complexity of alpha and beta, it would be interesting to see insights into the computational cost and the convergence behavior especially for language model training. \n- Questionable novelty. The key contributions of the paper closely resembles existing techniques that have explored using different divergences for distillation (similar to the alphanet paper, which the authors have cited). The proposed approach does not introduce a fundamentally novel concept. I would suggest to specify in the paper a more thorough comparison with prior work and clearly stating what would be the unique contribution. \n- Baselines in the experiments. I encourage the authors to add more KD baselines like the ones referenced above.", "other_comments_or_suggestions": "Please refer to previous cell.", "questions_for_authors": "Please refer to the strengths and weaknesses section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A4xAm1D6Gs", "reviewer": "Reviewer_Ty7Q", "summary": "The paper introduces ABKD, a knowledge distillation (KD) framework using alpha-beta-divergence to balance the \"hardness-concentration\" (focus on high-error classes) and \"confidence-concentration\" (focus on high-confidence classes) effects. Theoretical analysis shows that FKLD and RKLD represent extreme cases of these effects, leading to suboptimal performance. ABKD generalizes these divergences, enabling smooth interpolation via alpha and beta.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comments and suggestions. Our response follows:\n> Q1: The experiments should include a comparison of the performance of ABKD in these degenerate cases with the original FKLD and RKLD\n\n**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.\n> Q2: Theorem 3.2: Informal statement lacks rigor; formal proof needs verification\n\n**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.\n> 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\n\n**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'.\n\nTo 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.\n\n||Dolly|Self-Instruct|Vicuna|Super-Natural|Unnatural|\n|-|-|-|-|-|-|\n|SFT|24.54|16.80|16.15|29.29|27.43|\n|FKLD|25.23|18.90|16.67|31.68|29.36|\n|RKLD|27.74|20.61|18.83|35.31|33.86|\n|Jensen's KL|26.28|18.84|17.81|30.92|29.79|\n|BDKD|26.78|18.94|17.81|32.15|30.89|\n|AKL|26.38|17.69|16.72|33.02| 31.29|\n|DISTILLM|28.24|21.00|19.12|37.06|35.05|\n|AlphaNet|28.11|21.30|18.70|37.86|35.40|\n|Ours|**30.25**|**22.39**|**20.83**|**38.51**|**38.66**|\n> Q4: The Need for Ablation on Model Architectures.\n\n**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.\n\n> Q5: Essential References Not Discussed\n\n**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.\nWe 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.\n\n||Dolly|Self-Instruct|Vicuna|Super-Natural|Unnatural|\n|-|-|-|-|-|-|\n|AKL|23.83|10.87|15.63|20.07|21.97|\n|Ours|**25.65**|**13.47**|**16.06**|**26.47**|**29.32**|\n> Q6: Baseline Variance: Some baselines (e.g., LSD, TTM) underperform in Fig.5; unclear if hyperparams were optimized.\n\n**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.\n> Q7: Cross-Architecture Experiment\n\n**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.\n||KD|ABKD (Ours)|DKD|ABDKD (Ours)|LSD|ABLSD (Ours)|TTM|ABTTM (Ours)|\n|-|-|-|-|-|-|-|-|-|\n|Accuracy|73.81|74.62|74.37|**75.26**|74.52|74.77|74.87|75.02|\n\n> **Q8**: How does ABKD perform with α/β outside [0,1] (e.g., α=1.5, β=-0.5)? \n\n**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:\n\n|α\\β|-0.1|-0.3|-0.5|\n|-|-|-|-|\n|1.2 |70.81|71.10|70.35 |\n|1.4|71.29| 71.24| 70.92|\n|1.6|70.55|70.53|70.34|\n\nAn overly large α weakens the hardness-concentration and an overly small β weakens the confidence-concentration effect, both may degrade distillation performance.\n> Q9: Computational Cost Analysis\n\n**A9**: We compared the training costs of different methods on language modeling, as shown below.\n\n||Training Cost (second/sample)|\n|-|-|\n|SFT|0.344|\n|KD|0.649|\n|MiniLLM|4.452|\n|GKD|2.078|\n|DISTILLM|1.331|\n|AlphaNet|0.882|\n|Ours|0.768|\n\nOur 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.", "claims_and_evidence": "Supported Claims:  \nBalancing mode-concentration effects (Sec. 3–4): Theoretical analysis (Prop. 3.1, 4.2) and visualizations (Fig. 1d–g) justify the limitations of FKLD/RKLD and ABKD’s trade-off mechanism. \nPotential Issues:\nThe paper claims that ABKD encompasses FKLD (when alpha=1, beta=0) and RKLD (when alpha=0, beta=1), and thus the experiments should include a comparison of the performance of ABKD in these degenerate cases with the original FKLD and RKLD.", "methods_and_evaluation_criteria": "17 datasets and 12 teacher-student pairs (e.g., GPT-2 XL→GPT-2) are reasonable. Metrics (ROUGE-L, accuracy) align with standard practice. \nHowever, the validation on the CIFAR100 dataset only involves cases where the teacher and student have the same architecture; the performance when they have different architectures should also be examined.", "theoretical_claims": "Theorem 3.2(mass allocation differences): Informal statement lacks rigor; formal proof needs verification.", "experimental_designs_or_analyses": "Hyperparameter Sensitivity: Fig. 6 shows alpha and beta impact entropy and Self-BLEU, but optimal values vary across tasks (e.g., alpha=0.2, beta=0.7 for NLP vs. alpha=0.6, beta=0.5 for CIFAR-100 in Tabs. 4,6). This suggests task-specific tuning, weakening the \"universal\" claim.  \nMoreover, Table 4 in Appendix I.2.3 also shows that ABKD requires specific hyperparameter selection, which limits its scalability to more datasets and architectures.", "relation_to_broader_scientific_literature": "Builds on classical KD (Hinton, 2015) and recent divergence variants (MiniLLM, DISTILLM). The alpha-beta-divergence generalizes prior work (Table 1), addressing gaps in FKLD/RKLD trade-offs.", "essential_references_not_discussed": "Wu, T., Tao, C., Wang, J., Zhao, Z., & Wong, N. (2024). Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models. arXiv.Org. https://doi.org/10.48550/arxiv.2404.02657\nCui, X., Qin, Y., Gao, Y., Zhang, E., Xu, Z., Wu, T., Li, K., Sun, X., Zhou, W., & Li, H. (2024). SinKD: Sinkhorn Distance Minimization for Knowledge Distillation. IEEE Transactions on Neural Networks and Learning Systems, 1–15. https://doi.org/10.1109/tnnls.2024.3501335\nThese papers are representative works that improve distillation effects based on different loss distance calculations, but they are not cited in this paper. The adaptive KL divergence(AKL) proposed in the former is more suitable than the WSD proposed in this paper as a baseline.", "other_strengths_and_weaknesses": "Strengths:  \nTheoretical-empirical synergy: Clear connection between gradient analysis (Sec. 3) and empirical results.  \nWeaknesses:  \nBaseline Variance: Some baselines (e.g., LSD, TTM) underperform in Fig. 5; unclear if hyperparameters were optimized.", "other_comments_or_suggestions": "If you have any other comments or suggestions (e.g., a list of typos), please write them here. Clarity: The paper is well-structured, but theoretical sections (Sec. 3–4) are dense.", "questions_for_authors": "Q1: How does ABKD perform with alpha/beta outside [0,1] (e.g., alpha=1.5, beta=-0.5)? Does the framework still hold? (Clarifies generality claims.)\nQ2: Were computational costs (e.g., GPU hours) comparable between ABKD and baselines? (Addresses scalability concerns.)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wWJKgoeq2M", "reviewer": "Reviewer_VNjn", "summary": "This paper investigates a fundamental challenge in Knowledge Distillation (KD): the improper allocation of probability mass when using traditional divergences like Forward KL Divergence (FKLD) and Reverse KL Divergence (RKLD). FKLD tends to spread probability mass too broadly, failing to pay sufficient attention to the target class, while RKLD overly concentrates on the target class, neglecting the broader distributional information from the teacher model. The authors analyze this issue from the view of two effects:  hardness-concentration, which focuses on modes where the student model has high error, and confidence-concentration, which emphasizes modes where the student is already confident. To better balance the two effects, the paper introduces $\\alpha$-$\\beta$-Divergence (ABKD), which generalizes FKLD, RKLD, and other divergences, thus providing a better trade-off between the two effects. Theoretical results demonstrate that ABKD provides a more balanced allocation of probability mass, leading to improved student learning. Finally, extensive experiments across 17 language and vision datasets with 12 teacher-student model pairs validate its effectiveness.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> 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\n\n**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).\n\n|GPT-2 XL (1.5B) -> GPT-2 (0.1B, 0.3B, 0.8B) | Dolly Eval | Self-Instruct|Vicuna Eval|Super-Natural|Unnatural|\n|-|-|-|-|-|-|\n| **α**  | 0.2        | 0.2         | 0.2        | 0.2          | 0.2         | \n| **β**  | 0.7        | 0.7         | 0.7        | 0.7          | 0.7         | \n\n\n| **Dataset**      | ImageNet   | Caltech101  | OxfordPets | StanfordCars | Flowers102  | Food101    | FGVCAircraft | SUN397   | DTD   | EuroSAT | UCF101  |\n|-|-|-|-|-|-|-|-|-|-|-|-|\n| **α**  | 0.5        | 0.8         | 0.8        | 0.6          | 0.9         | 0.5        | 0.6          | 0.8      | 1.0   | 0.6     | 0.8     |\n| **β**  | 0.5        | 0.2         | 0.4        | 0.4          | 0.1         | 0.5        | 0.5          | 0.2      | 0.2   | 0.5     | 0.2     |\n\nTo 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.\n\nBased 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). \n\nThe results below show that our method outperforms others by 0.65-3.26, especially excelling in Dolly and Unnatural.\n\n| Method    | Dolly Eval | Self-Instruct | Vicuna Eval | Super-Natural | Unnatural |\n|-----------|-----------|---------------|-------------|---------------|-----------|\n| SFT       | 24.54 (0.51) | 16.80 (0.64) | 16.15 (0.15) | 29.29 (0.13) | 27.43 (0.21) |\n| FKLD        | 25.23 (0.44) | 18.90 (1.20) | 16.67 (0.35) | 31.68 (0.22) | 29.36 (0.13) |\n| RKLD   | 27.74 (0.45) | 20.61 (0.80) | 18.83 (0.40) | 35.31 (0.24) | 33.86 (0.16) |\n| Jensen's KL       | 26.28 (0.43) | 18.84 (0.66) | 17.81 (0.38) | 30.92 (0.12) | 29.79 (0.17) |\n| BDKD       | 26.78 (0.53) | 18.94 (0.68) | 17.81 (0.52) | 32.15 (0.34) | 30.89 (0.24)  |\n| AKL       | 26.38 (0.41) | 17.69 (0.46) | 16.72 (0.48) | 33.02 (0.16) | 31.29 (0.08) |\n| DISTILLM  | 28.24 (0.48) | 21.00 (0.72) | 19.12 (0.53) | 37.06 (0.35) | 35.05 (0.13) |\n| AlphaNet  | 28.11 (0.29) | 21.30 (0.63) | 18.70 (0.23) | 37.86 (0.44) | 35.40 (0.17) |\n| Ours (ABKD) | **30.25** (0.37) | **22.39** (0.62) | **20.83** (0.42) | **38.51** (0.32) | **38.66** (0.10) |", "claims_and_evidence": "The claims in the submission are well-supported by both theoretical analysis and empirical results:\n\n- Theoretical justifications explain the limitations of FKLD and RKLD and demonstrate how ABKD provides a more flexible probability mass allocation. \n- Empirical evaluations across 17 language and vision datasets with 12 teacher-student configurations further validate these claims, showing consistent performance improvements over existing methods.", "methods_and_evaluation_criteria": "Yes. The evaluation protocol is comprehensive, spanning 17 language/vision datasets with 12 teacher-student configurations, ensuring that the findings generalize across different settings. The selection of benchmarks, including instruction-following datasets for NLP and classification datasets for vision tasks, is appropriate for assessing the effectiveness of distillation methods. Moreover, the proposed method is compared against the state-of-the-art KD techniques, also providing meaningful results.", "theoretical_claims": "Yes. I have carefully checked the theoretical claims presented in the paper, particularly ABKD's role in balancing the two effects.", "experimental_designs_or_analyses": "Yes, please refer to `Methods And Evaluation Criteria`.", "relation_to_broader_scientific_literature": "Nowadays, knowledge distillation is an important topic due to the development of the large-scale foundation model. This paper investigates a foundational issue in the field of knowledge distillation. Hence, I believe it could relate to broader literature in the future.", "essential_references_not_discussed": "To my best of knowledge, this paper has covered essential references in this field.", "other_strengths_and_weaknesses": "The strengths of this paper are summarized as follows:  \n\n- **Rigorous theoretical insights**: This paper provides a rigorous theoretical analysis of the limitations of traditional divergences widely used in KD. The key idea behind is to reveal how these divergences allocate probability mass during training. By introducing the concepts of hardness-concentration and confidence-concentration, the authors show that FKLD and RKLD represent two extreme cases: FKLD fails to concentrate on the target class and RKLD overly concentrates on the target class. The proposed $\\alpha$-$\\beta$-Divergence (ABKD) balances the two effects, offering a novel perspective on why model-based guidance (i.e., distilling from a teacher model) can outperform hard label guidance (i.e., learning from one-hot labels).\n\n- **Flexibile framework based on $\\alpha$-$\\beta$-divergence**: The proposed framework, based on $\\alpha$-$\\beta$-divergence, unifies existing divergence-based KD approaches. This flexibility allows for fine-grained control over probability mass allocation, leading to improved performance.  \n\n- **Comprehensive empirical validation**: Extensive experiments are conducted across 17 language/vision datasets with 12 teacher-student model configurations. The results demonstrate consistent improvements over state-of-the-art KD methods.\n\nMy minor concerns are as follows:  \n\n- In Sec 5.3, the authors provide empirical justifications for the choice of $\\alpha$ and $\\beta$. Although the results are valuable for hyperparameter tuning, 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.", "other_comments_or_suggestions": "Please refer to `weankess`.", "questions_for_authors": "Please refer to `weankess`.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Reviewers were unanimously positive in their assessment of the paper, which studies popular KL and reverse-KL based distillation objectives through a unified framework. The framework highlights the importance of Hardness- and Confidence-Concentration, which motivates a proposed $(\alpha,\beta)$-divergence based objective. The paper demonstrates results on a range of language and vision settings, with several model families.

The author response included some new results comparing the method to a further set of baselines, including the Jensen-Shannon divergence. These results are strongly encouraged to feature in the final submission.

Overall, the paper is well written, studies an important problem, and makes interesting contributions that are likely of interest to the community. Thus, we believe it is worthy of presentation at the conference.
Decision: Accept (oral)