Research map / Label-efficient and robust learning / Knowledge distillation
ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via $\alpha$-$\beta$-Divergence
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ABKD, a knowledge distillation objective based on the alpha-beta-divergence that interpolates between forward and reverse KL to balance hardness- and confidence-concentration effects, evaluated on instruction-following (GPT-2 family) and CIFAR-100/vision distillation.
Problem
Forward KL yields overly smooth student distributions that fail to focus on the target class, while reverse KL collapses toward one-hot supervision; the paper seeks a divergence that properly allocates probability mass in knowledge distillation.
Approach
Introduce the alpha-beta-divergence as a two-parameter family that unifies FKLD (alpha=1, beta=0) and RKLD (alpha=0, beta=1) and recovers alpha-, beta-, Hellinger and squared Euclidean divergences, letting alpha control hardness-concentration and beta control confidence-concentration independently.
Key results
- On five instruction-following datasets distilling GPT-2 XL (1.5B) into GPT-2 (0.1B), ABKD improves over FKLD and RKLD by 0.81 to 3.31 ROUGE-L points.
- ABKD achieves 29.92 ROUGE-L on Unnatural for GPT-2 XL to GPT-2 Large, the best among compared methods in that setting.
- ABKD matches vanilla KD training speed, whereas SGO-based methods require 1.6 to 7 times longer training time.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace FKLD distillation loss with the α-β-divergence in ABKD | GPT-2 XL → GPT-2 · Dolly Eval | ROUGE-L: 25.65 (0.24) vs 24.62 (0.33) | Helped |
| Replace FKLD distillation loss with the α-β-divergence in ABKD | GPT-2 XL → GPT-2 · Self-Instruct | ROUGE-L: 13.47 (0.42) vs 11.65 (0.28) | Helped |
| Replace FKLD distillation loss with the α-β-divergence in ABKD | GPT-2 XL → GPT-2 · Vicuna Eval | ROUGE-L: 16.06 (0.25) vs 17.30 (0.41) | Hurt |
| Replace FKLD distillation loss with the α-β-divergence in ABKD | GPT-2 XL → GPT-2 · Super-Natural | ROUGE-L: 26.47 (0.31) vs 23.76 (0.38) | Helped |
| Replace FKLD distillation loss with the α-β-divergence in ABKD | GPT-2 XL → GPT-2 · Unnatural | ROUGE-L: 29.32 (0.08) vs 25.79 (0.08) | Helped |
| Constrain the divergence to α-divergence rather than independently tuning α and β (ablation) | ABKD · Instruction-following datasets | ROUGE-L: 0.81 to 3.31 | Hurt |
| Use weighted-sum divergence (WSD), combining weighted FKLD and RKLD (ablation) | ABKD · Instruction-following datasets | ROUGE-L | Hurt |
| Decrease α to strengthen hardness-concentration (ablation) | ABKD · CIFAR-100 | accuracy | No clear effect |
| Decrease α to strengthen hardness-concentration (ablation) | ABKD · Dolly | ROUGE-L | Helped |
| Decrease β to weaken confidence-concentration (ablation) | ABKD · Instruction generation | Self-BLEU | Helped |
Surprising findings
- Alpha-divergence alone (constrained to alpha+beta=1) underperforms, and weighted-sum divergence fails to match the alpha-beta-divergence, suggesting unconstrained two-parameter control matters.
Limitations
- Hyperparameters alpha and beta require tuning, with no universal setting; grid search adds overhead.
- Instruction-following comparisons use fixed datasets for ABKD versus augmented data for other methods, making the main comparison unfair by the authors' admission.
Experimental setup
Datasets: Dolly Self-Instruct Vicuna Eval Super-Natural Instructions Unnatural Instructions CIFAR-100
Benchmarks: Dolly Eval Self-Instruct Vicuna Eval Super-Natural Instructions Unnatural Instructions CIFAR-100
Metrics: ROUGE-L accuracy harmonic mean (base-to-new) Self-BLEU Shannon entropy
Baselines: SFT KD SeqKD MiniLLM GKD DISTILLM DKD LSD TTM CoCoOp MaPLe PromptSRC
Builds on: forward KL distillation reverse KL distillation (MiniLLM) alpha-divergence DKD GKD
Closest papers by meaning
- Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models (NeurIPS 2026 · similarity 0.91)
- Why Knowledge Distillation Works in Generative Models: A Minimal Working Explanation (NeurIPS 2025 · similarity 0.90)
- DeepKD: A Deeply Decoupled and Denoised Knowledge Distillation Trainer (NeurIPS 2025 · similarity 0.90)
- Enhancing Logits Distillation with Plug&Play Kendall's $\tau$ Ranking Loss (ICML 2025 · similarity 0.89)
- PLD: A Choice-Theoretic List-Wise Knowledge Distillation (NeurIPS 2025 · similarity 0.89)
- TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.