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

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

Speakers: Guanghui Wang, Zhiyong Yang, Zitai Wang, Shi Wang, Qianqian Xu, Qingming Huang

Conference: ICML 2025

YouTube: https://slideslive.com/39044066

Overview

This talk introduces ABKD (Alpha-Beta Knowledge Distillation), a novel framework designed to enhance model compression through knowledge distillation by addressing the limitations of traditional divergence measures. As foundation models continue their exponential growth, driven by scaling laws that correlate performance with model size, the computational and financial demands for deploying and fine-tuning these models have become prohibitive for most users. Knowledge Distillation (KD) stands out as a highly effective model compression technique, widely adopted in the development of popular foundation models like DeepSeek and Qwen to achieve significant performance gains in smaller models.

The core problem ABKD tackles lies in the distribution matching component of knowledge distillation, where the goal is to align the output probability distributions of a large, high-performing "teacher" model with a smaller "student" model. Existing KD approaches predominantly rely on Kullback-Leibler (KL) divergence, which, due to its asymmetry, can be applied in two directions: Forward KL (FKLD) or Reverse KL (RKLD). However, the choice between these has lacked systematic theoretical analysis, leading to suboptimal or inconsistent performance. This paper provides a rigorous theoretical examination of the learning dynamics induced by FKLD and RKLD, revealing their extreme and often undesirable effects on how the student model reallocates probability across classes.

To overcome these limitations, the authors propose a more flexible and robust solution based on alpha-beta divergence, a two-parameter family of divergence measures that can smoothly interpolate between and generalize beyond FKLD and RKLD. ABKD offers independent control over the student's learning dynamics, allowing for tailored probability allocation that can prevent issues like conservative learning, aggressive overfitting to a potentially flawed teacher, loss explosion, and vanishing gradients. The framework's efficacy is demonstrated through extensive experiments across diverse NLP and vision tasks, showcasing consistent performance improvements over vanilla KD and supervised fine-tuning (SFT) without incurring additional training overhead.

Background

▶ Watch: Introduction to Knowledge Distillation and problem statement (0:00)

The current landscape of artificial intelligence is dominated by foundation models, which have demonstrated unprecedented capabilities across a wide array of tasks. A key characteristic of these models is their adherence to scaling laws, where performance tends to improve predictably with increases in model size, data, and computational resources. This has led to an relentless push towards ever-larger models, exemplified by architectures like DeepSeek and Qwen, which often boast billions or even trillions of parameters. While these colossal models achieve state-of-the-art results, their sheer scale translates into immense computational costs for both training and inference, making them inaccessible for many applications and users. The demand for efficient model compression techniques has thus become paramount.

Knowledge Distillation (KD) has emerged as a leading strategy for compressing large teacher models into smaller, more efficient student models while retaining much of the teacher's performance. At its heart, KD involves training a compact student model to mimic the behavior of a larger, pre-trained teacher model. A common and effective approach within KD is distribution matching, where the student's output class probabilities are encouraged to align with those of the teacher. The objective function for this typically comprises two parts: a standard cross-entropy loss on the ground truth labels and a divergence term that quantifies the dissimilarity between the teacher's probability distribution (P) and the student's probability distribution (Q). This divergence term is crucial as it provides the student with rich, fine-grained supervision beyond just the hard labels, guiding it to shape its entire output distribution.

Historically, the choice of divergence measure in KD has largely revolved around the Kullback-Leibler (KL) divergence. However, KL divergence is inherently asymmetric. This asymmetry means that the direction of computation—either from teacher to student (Forward KL, P || Q) or from student to teacher (Reverse KL, Q || P)—yields fundamentally different learning dynamics. Despite its widespread use, a systematic theoretical analysis of which KL divergence direction is optimal and why has been largely absent. Empirical observations often show that FKLD can be too "conservative" or "lazy," slow to adapt and change the student's distribution, while RKLD can be too "aggressive," causing the student to blindly follow the teacher's output, even if the teacher might be incorrect or overly confident in certain regions. This lack of a principled understanding and a flexible mechanism to balance these opposing dynamics highlights a significant gap in the effective application of knowledge distillation.

Key Findings

▶ Watch: One-parameter interpolation with Alpha Divergence (3:30)

The research presented in this talk makes several pivotal contributions to the field of knowledge distillation, primarily by providing a rigorous theoretical foundation and a practical framework to overcome the limitations of existing divergence measures.

1. Systematic Analysis of KL Divergence Dynamics: The authors provide the first systematic theoretical analysis of how Forward KL Divergence (FKLD) and Reverse KL Divergence (RKLD) influence the student model's probability reallocation across classes. Through the use of a log R term (analogous to log-likelihood ratio in RLHF), they decompose the learning dynamics into two key effects: a confidence weighting term based on the student's distribution (Q), and a matching error term reflecting the disagreement between teacher (P) and student (Q). This analysis reveals that FKLD induces a "weak" or conservative dynamic, making the student "lazy" to change, while RKLD induces an "aggressive" dynamic, causing the student to fit the teacher's guidance too strongly, even when the teacher might be suboptimal. This imbalance highlights the need for a more nuanced approach.

1. Introduction of Alpha-Beta Divergence for Flexible Control: To strike a proper balance, the paper proposes leveraging alpha-beta divergence, a two-parameter family of f-divergences. While single-parameter divergences like alpha divergence can interpolate between FKLD and RKLD, they suffer from a constraint where improving one effect (e.g., confidence weighting) necessitates reducing another (e.g., hardness concentration), leading to "improper probability allocation." Alpha-beta divergence breaks this constraint by offering independent control over the confidence and hardness concentration effects via its two parameters, alpha and beta. This provides a significantly more general and flexible framework for designing KD objectives.

1. The ABKD Framework: Based on these theoretical insights, the authors introduce ABKD (Alpha-Beta Knowledge Distillation). This framework directly incorporates alpha-beta divergence into the knowledge distillation objective function. ABKD not only unifies and recovers classical KL divergences as special cases but also generalizes to a broader family of divergence measures previously unexplored in KD, such as Hellinger distance and log Euclidean divergence. This continuous surface of divergences allows for fine-grained tuning of the student's learning behavior.

1. Robustness and Gradient Stability: A crucial practical finding is that ABKD offers enhanced robustness compared to FKLD or RKLD. It is significantly more stable in extreme scenarios, such as when either the teacher's (P) or student's (Q) probability for a class is very small, which can lead to loss explosions and gradient instability with traditional KL divergences. Furthermore, ABKD helps to prevent gradient vanishing when P and Q are far apart, a common occurrence in the early stages of training due to the inherent capacity gap between teacher and student models. This stability ensures more reliable and faster convergence.

1. Empirical Validation and Performance Gains: Through extensive experiments across a variety of tasks, including instruction following, mathematical reasoning (NLP), and vision tasks, ABKD consistently demonstrates superior performance. When distilling large models like GPT-2 XL into smaller GPT-2 variants, Open Llama 2 (7B) into 3B, and Qwen 2.5 (7B) into 1.5B, ABKD consistently outperforms both vanilla KD and Supervised Fine-Tuning (SFT) baselines. Notably, ABKD achieves these gains without introducing any additional computational overhead during training, as it only modifies the divergence objective. For instance, in mathematical reasoning, ABKD achieves an average improvement of 1.1 over FKLD across various benchmarks.

Technical Deep Dive

▶ Watch: Introducing Alpha-Beta Divergence for independent control (4:40)

The technical core of this work lies in its re-evaluation of divergence measures within knowledge distillation and the introduction of a more generalized framework.

In knowledge distillation, the student model (Q) is trained to mimic the teacher model (P). The overall objective function typically combines a standard cross-entropy loss with a divergence term:

$$L_{KD} = L_{CE}(y, Q) + \lambda \cdot D(P || Q)$$

where $L_{CE}$ is the cross-entropy loss with ground truth labels $y$, and $D(P || Q)$ is the divergence between teacher distribution P and student distribution Q, scaled by $\lambda$.

The talk highlights the limitations of the two most common choices for $D(P || Q)$:

1. Forward KL Divergence (FKLD): $D_{KL}(P || Q) = \sum_i P_i \log \frac{P_i}{Q_i}$. This formulation penalizes the student heavily when $Q_i$ is small but $P_i$ is large. It encourages the student to cover all modes of the teacher's distribution. The analysis shows FKLD to be "conservative" or "lazy," meaning the student's probability reallocation (quantified by the log R term) is weakly influenced. It emphasizes regions where the student is confident (high Q), but the matching error term is also weighted by Q, making it less aggressive in regions of disagreement.
2. Reverse KL Divergence (RKLD): $D_{KL}(Q || P) = \sum_i Q_i \log \frac{Q_i}{P_i}$. This formulation penalizes the student when $Q_i$ is large but $P_i$ is small, essentially preventing the student from predicting classes that the teacher deems unlikely. It encourages the student to be confident only on modes where the teacher is also confident. RKLD is characterized as "aggressive," as it strongly drives the student to fit the teacher's guidance. The log R term here shows that the confidence weighting is weak (weighted by P, which might be small), but the matching error term is strongly weighted by Q, pushing the student to align strongly with the teacher's peaks. The core issue is that neither FKLD nor RKLD provides a "proper" rate of learning, often being either too slow or too susceptible to the teacher's potential errors.

To overcome this, the paper first explores alpha divergence, a one-parameter family of divergences defined as a geometric average of P and Q (though the exact formula isn't explicitly shown, it's a known family). As the parameter $\alpha$ approaches 1, it converges to FKLD; as $\alpha$ approaches 0, it converges to RKLD. While alpha divergence offers an interpolation, its learning dynamics (specifically the log R term) show that the powers controlling confidence and hardness concentration effects must add up to one. This constraint means that increasing the strength of one effect necessarily reduces the other, limiting the flexibility in probability reallocation.

The key innovation is the adoption of alpha-beta divergence, a two-parameter generalization that removes this constraint. The talk indicates its definition (which is typically $D_{\alpha, \beta}(P || Q) = \frac{1}{\alpha \beta} \left( \sum_i P_i^\alpha Q_i^\beta - 1 \right)$ for certain ranges, or a more general f-divergence form). The critical theoretical result for alpha-beta divergence is that its log R term now features two independent power parameters, $\alpha$ and $\beta$. This independence is crucial because it allows for separate control over:

• Confidence weighting: How much the student's own confidence (Q) influences its learning.
• Hardness concentration: How much the disagreement between P and Q drives the student to adapt.

By varying $\alpha$ and $\beta$ independently, the ABKD framework can construct a continuous surface of divergence measures, where FKLD and RKLD represent extreme points, and alpha divergence is a sub-manifold. This flexibility allows ABKD to recover not only KL divergences but also other known divergences like Hellinger distance and log Euclidean divergence as special cases, thus unifying a broad family of measures under a single framework.

Beyond flexible control, ABKD offers enhanced numerical stability. Traditional KL divergences can cause the loss to explode when probabilities ($P_i$ or $Q_i$) become extremely small, leading to noisy and unstable gradients. ABKD's formulation is more robust to these extreme modes. Furthermore, in early training stages, when the teacher and student models have a significant capacity gap, their distributions P and Q can be far apart. In such scenarios, divergences like Jensen-Shannon (JS) divergence can degenerate into a constant, leading to vanishing gradients and hindering convergence. ABKD's structure inherently mitigates this issue, ensuring more stable gradient flow and faster convergence, particularly when $P$ and $Q$ are sufficiently different.

The talk provides theoretical guidance for hyperparameter tuning of $\alpha$ and $\beta$:

• A smaller $\alpha$ facilitates escaping local minima and leads to faster convergence, especially when the teacher and student distributions are initially divergent (common in early training).
• A smaller $\beta$ encourages the student to absorb knowledge from the teacher regarding non-target classes, which is crucial for improving generalization.

This principled understanding transforms the selection of a divergence function from an empirical guessing game into a theoretically informed decision, tailored to the specific learning dynamics desired.

Experimental Setup & Results

▶ Watch: ABKD's robustness and gradient stability advantages (6:30)

The effectiveness of the ABKD framework was rigorously validated through a comprehensive series of experiments spanning both natural language processing (NLP) and computer vision (CV) tasks. The experiments focused on demonstrating ABKD's superior performance compared to existing distillation methods and its robustness under various conditions.

Hyperparameter Tuning Guidance:

Before diving into specific task results, the authors established general guidelines for tuning $\alpha$ and $\beta$:

• Smaller $\alpha$: Generally speeds up convergence, particularly useful when teacher and student distributions are significantly different, helping the student escape local minima.
• Smaller $\beta$: Leads to a smoother student distribution, encouraging better generalization by absorbing knowledge about non-target classes.
• Simple Classification Tasks: To prevent overfitting, a larger $\alpha$ (to slow down convergence) and a smaller $\beta$ (for smooth distribution) were found to be effective.
• Complicated Tasks (e.g., Instruction Following with many classes): Here, the goal is often to speed up training and focus on target classes. This typically calls for a smaller $\alpha$ and a larger $\beta$. This highlights the task-dependent nature of optimal $\alpha, \beta$ settings, but with clear theoretical underpinnings.

Instruction Following Task (NLP):

• Teacher Model: GPT-2 XL (a large generative pre-trained transformer).
• Student Models: Smaller GPT-2 variants.
• Datasets: A range of datasets designed for instruction following.
• Baselines: Vanilla KD (presumably using FKLD or RKLD), and Supervised Fine-Tuning (SFT).
• Key Results:
• Distillation methods generally outperformed SFT, underscoring their potential for transfer learning and model compression.
• ABKD consistently outperformed both vanilla KD and SFT across the entire range of datasets used for distilling GPT-2 XL into smaller GPT-2 models. This demonstrates the framework's ability to provide a more effective distillation objective.
• An important observation was that KD could sometimes underperform SFT, for example, on "unnatural" datasets when distilling GPT-2 XL into GPT-2. This emphasizes the critical importance of choosing the right distillation objective, which ABKD aims to provide.

Large Language Model (LLM) Distillation:

• Teacher Model: Open Llama 2 (7B parameters).
• Student Model: A 3B parameter variant of Open Llama 2.
• Benchmarks: Various standard LLM benchmarks (specific names not detailed in the transcript, but implied to be comprehensive).
• Key Results: ABKD consistently outperformed state-of-the-art baselines, showcasing its efficacy in compressing large, modern LLMs while retaining performance.

Mathematical Reasoning Task (NLP):

• Teacher Model: Qwen 2.5 (7B parameters).
• Student Model: A 1.5B parameter variant of Qwen 2.5.
• Benchmarks: Multiple mathematical reasoning benchmarks.
• Key Results: With a simple calibration of FKLD (implying a baseline comparison), ABKD achieved an average improvement of 1.1 across various mathematical reasoning benchmarks. This suggests a significant boost in the student model's reasoning capabilities compared to traditional KD.

Vision Tasks:

• The talk briefly mentions that the framework's effectiveness was also validated on "vision tasks," though no specific models, datasets, or quantitative results are provided in the transcript. This suggests the generality of ABKD beyond just NLP.

Training Cost Analysis:

• A crucial practical aspect addressed was the computational overhead. Figure 3a (referenced in the talk) indicated that ABKD trains just as fast as vanilla KD. This is because ABKD only modifies the distillation objective function itself, without adding any extra layers, complex architectures, or additional computational steps during the forward or backward passes. This addresses concerns about scalability, making ABKD a practically viable solution for large-scale model compression.

In summary, the experimental results robustly support the theoretical claims of ABKD, demonstrating its ability to yield more performant and stable student models across diverse domains and model sizes, all while maintaining training efficiency.

Practical Implications

▶ Watch: Practical guidance for tuning Alpha and Beta (7:50)

The ABKD framework carries significant practical implications for various stakeholders involved in the development and deployment of AI/ML models, particularly in the era of large foundation models.

For Practitioners and Model Builders:

• Principled Divergence Selection: ABKD provides a principled, theoretically grounded method for choosing the divergence measure in knowledge distillation, moving beyond the often ad-hoc or empirical choices of FKLD and RKLD. Model builders can now make informed decisions about the learning dynamics they want to impart to the student model.
• Fine-Grained Control over Learning: The independent control offered by the $\alpha$ and $\beta$ parameters allows for fine-tuning the student's learning behavior. This means practitioners can explicitly balance how much the student relies on its own confidence versus how aggressively it tries to match the teacher's distribution, or how much it focuses on target classes versus absorbing knowledge on non-target classes. This level of control is invaluable for optimizing performance on specific tasks or datasets.
• Robustness and Stability: The framework's enhanced robustness to extreme probability values (where P or Q are very small) and its ability to prevent vanishing gradients in early training are major advantages. This translates to more stable training, faster convergence, and reduced debugging time, especially when dealing with complex, high-dimensional probability distributions typical of large models.
• Improved Generalization: By allowing for smoother distributions (smaller $\beta$) and better absorption of teacher knowledge on non-target classes, ABKD can lead to student models with superior generalization capabilities, which is a critical factor for real-world deployment.

For Infrastructure Teams and Deployers:

• Efficient Model Compression: ABKD directly facilitates the creation of smaller, more efficient student models that retain a higher percentage of the teacher's performance. This is crucial for reducing inference costs, lowering memory footprints, and enabling deployment on edge devices or environments with limited computational resources. The ability to distill a 7B parameter model into a 3B model or a 7B into a 1.5B model with minimal performance loss has direct economic and operational benefits.
• Scalability: The finding that ABKD trains just as fast as vanilla KD, incurring no additional computational overhead, is vital for scalability. Infrastructure teams can integrate ABKD into existing training pipelines without needing to provision additional compute resources or extend training schedules.

Tradeoffs and Limitations:

• Increased Hyperparameter Tuning: While the theoretical guidance for $\alpha$ and $\beta$ is a significant step forward, it still introduces two additional hyperparameters that need to be tuned. While general guidelines are provided (e.g., smaller $\alpha$ for faster convergence, smaller $\beta$ for smoother distributions), finding the absolute optimal values for a specific task might still require some experimentation or a systematic search strategy. This adds a layer of complexity compared to simply choosing FKLD or RKLD.
• Understanding Dynamics: Effectively utilizing ABKD requires a deeper understanding of the desired learning dynamics. Practitioners need to conceptualize what "faster convergence" or "smoother distribution" means for their specific model and task to appropriately set $\alpha$ and $\beta$.
• Teacher Quality Assumption: Like all KD methods, ABKD assumes a "good" teacher model. If the teacher itself is flawed or biased, the student will still learn from those imperfections, albeit with more controlled dynamics. The framework provides tools to prevent blindly following a potentially wrong teacher (by adjusting $\alpha, \beta$ to be less aggressive), but it doesn't fundamentally fix a poor teacher.

Overall, ABKD offers a powerful and flexible tool for advancing knowledge distillation, empowering practitioners to build more efficient, robust, and high-performing compressed models for the demanding applications of today's AI landscape.

Key Takeaways

• The exponential growth of foundation models necessitates efficient compression techniques like Knowledge Distillation (KD) to make them affordable and deployable.
• Traditional KD, relying on Forward KL (FKLD) or Reverse KL (RKLD) divergence, suffers from inherent limitations: FKLD is too conservative, while RKLD is too aggressive, leading to suboptimal probability reallocation during student training.
• The alpha-beta divergence offers a flexible, two-parameter interpolation mechanism that provides independent control over the student's confidence weighting and hardness concentration effects, resolving the limitations of single-parameter divergences.
• The proposed ABKD (Alpha-Beta Knowledge Distillation) framework incorporates this divergence, unifying and generalizing classical KL divergences and other measures (e.g., Hellinger distance), leading to a more robust and stable KD objective.
• ABKD significantly improves training stability by preventing loss explosions in extreme probability scenarios and mitigating gradient vanishing when teacher and student distributions are initially far apart, ensuring faster and more reliable convergence.
• Empirical results show ABKD consistently outperforms vanilla KD and Supervised Fine-Tuning (SFT) across diverse NLP (instruction following, LLM distillation, mathematical reasoning) and vision tasks, without introducing any additional training computational overhead.
• Theoretical insights provide practical guidance for hyperparameter tuning: a smaller $\alpha$ facilitates faster convergence and escaping local minima, while a smaller $\beta$ encourages smoother distributions and better generalization.

About the Speaker(s)

The work presented in this talk was authored by Guanghui Wang, Zhiyong Yang, Zitai Wang, Shi Wang, Qianqian Xu, and Qingming Huang.