DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs

DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs

Speakers: Jongwoo Ko, KAIST AI; Tianyi Chen, Microsoft; Sungnyun Kim, KAIST AI; Tianyu Ding, Microsoft; Luming Liang, Microsoft; Ilya Zharkov, Microsoft; Se-Young Yun, KAIST AI

Conference: ICML 2025

YouTube: https://slideslive.com/39044065

Overview

In an era increasingly defined by the capabilities of large language models (LLMs), the computational demands associated with their deployment remain a significant bottleneck. This talk introduces DistiLLM-2, a novel knowledge distillation (KD) framework designed to address this challenge by enabling smaller, more efficient language models (SLMs) to achieve performance levels comparable to their larger counterparts. Presented by Jongwoo Ko from KAIST AI, representing a collaborative effort with Microsoft, DistiLLM-2 stands out as the first work to jointly optimize both the loss formulation and data creation aspects of knowledge distillation for LLMs.

The core motivation behind DistiLLM-2 stems from the recognition that while LLMs demonstrate unparalleled effectiveness, their extensive computational and memory requirements often preclude their use in resource-constrained environments, such as on-device applications. SLMs, being more compact and efficient, are ideally suited for such scenarios, but traditionally lack the advanced capabilities of LLMs. Knowledge distillation offers a pathway to bridge this gap, allowing SLMs to learn from the output distributions of powerful teacher LLMs. DistiLLM-2 significantly advances this field by proposing a contrastive approach that refines how SLMs learn, leading to more effective and efficient knowledge transfer.

This research is particularly impactful because it tackles a critical limitation of prior KD efforts, which often focused on either refining loss functions or curating training data in isolation. By demonstrating the synergistic potential of jointly optimizing these two perspectives, DistiLLM-2 not only boosts the performance of student SLMs but also offers a more robust and broadly applicable distillation methodology. The proposed framework, built upon a sophisticated understanding of KL divergences and incorporating adaptive, curriculum-based training strategies, promises to accelerate the deployment of high-performing AI models across a wider spectrum of applications, from edge devices to enterprise-level solutions.

Background

▶ Watch: DistiLLM-2: Jointly optimizing KD loss and data (0:00)

The rapid advancements in large language models (LLMs) have revolutionized natural language processing, enabling breakthroughs in areas like content generation, summarization, and complex reasoning. However, the sheer scale of these models, often comprising billions or even trillions of parameters, translates into substantial computational costs for both training and inference. This resource intensiveness—demanding significant GPU memory, processing power, and energy—poses a formidable barrier to their widespread deployment, especially in environments with limited resources such as mobile devices, embedded systems, or edge computing infrastructure.

Small language models (SLMs) offer a compelling alternative for these resource-constrained settings due to their compact size and efficiency. Yet, SLMs inherently possess less knowledge and exhibit lower performance compared to their larger counterparts. Knowledge Distillation (KD) emerges as a critical technique to bridge this performance gap. In KD, a large, pre-trained LLM (the teacher model) guides the training of a smaller SLM (the student model). The student learns to mimic the output distribution, or "knowledge," of the teacher, thereby acquiring advanced capabilities without incurring the teacher's full computational overhead.

Prior research in KD for language models has largely concentrated on two distinct avenues: either developing novel loss formulations to better capture the teacher's knowledge, or devising strategies for data creation to generate high-quality, relevant training examples. However, a significant gap identified by the DistiLLM-2 team was the lack of investigation into the synergy between these two crucial aspects. Most existing methods treated loss functions and data generation as separate concerns, potentially limiting the ultimate performance achievable by student models.

Simultaneously, the field of LLM preference optimization has seen the rise of contrastive approaches, such as Direct Preference Optimization (DPO) and Identity Preference Optimization (IPO). These methods go beyond merely increasing the likelihood of preferred responses; they actively decrease the likelihood of dis-preferred responses. This "contrastive" philosophy, which explicitly distinguishes between desired and undesired outcomes, has proven highly effective in aligning LLMs with human preferences. The DistiLLM-2 team recognized the potential of applying this powerful contrastive concept to the domain of knowledge distillation.

However, a direct application of DPO to KD presented significant challenges. Early attempts, like DPKD (Direct Preference KD), which replaced the reference model with a teacher and the policy model with a student, often led to an undesirable outcome for small SLMs: instead of effectively learning from the teacher's distribution, the student models tended to lose their pre-trained information. This observation was crucial, highlighting that a nuanced adaptation of contrastive principles was required for KD. The team realized that rather than directly manipulating likelihoods as in DPO, leveraging the inherent behaviors of forward KL (skew-KL) and reverse KL (skew-reverse-KL) divergences could provide a more stable and effective pathway for contrastive knowledge transfer. This foundational insight paved the way for the development of KALD and subsequently, DistiLLM-2, by implicitly controlling desired and undesired responses through these divergence measures.

Key Findings

▶ Watch: Rethinking contrastive KD: Leveraging KL divergences (4:00)

DistiLLM-2 represents a significant advancement in the field of knowledge distillation for large language models, primarily through its novel contrastive approach and the joint optimization of data creation and loss formulation. The key findings and contributions can be summarized as follows:

1. Introduction of KALD (Contrastive Approach for LLM Distillation): The research successfully introduced KALD, a framework that integrates skew-KL (forward KL) and skew-reverse-KL (reverse KL) divergences to create a powerful contrastive distillation mechanism. Unlike direct applications of DPO, KALD implicitly controls the student's learning by leveraging the distinct behaviors of these KL divergences: skew-KL encourages matching high-probability teacher outputs, while skew-reverse-KL discourages divergence from low-probability teacher outputs. This "response-aware" loss formulation led to faster convergence and superior performance compared to using either divergence alone.

1. Optimized Data Creation Strategy: Through extensive experimentation, the team identified the most effective strategy for data creation within the KALD framework. It was found that using Teacher-Generated Output (TGO) for the skew-KL component (representing desired responses) and Student-Generated Output (SGO) for the skew-reverse-KL component (representing undesired responses) consistently yielded the best results. This straightforward yet highly effective strategy optimizes the data used to guide the contrastive learning process. Furthermore, the option to use high-quality responses from superior LLMs (even beyond the immediate teacher) for TGO was explored, enhancing the quality of the "desired" signal.

1. Curriculum-Based Adaptive Loss Formulation (DistiLLM-2 Specific): Building upon KALD, DistiLLM-2 introduced two critical refinements to the loss function, addressing previous limitations and further boosting performance:

• Automated Interpolation Coefficients: A limitation in prior skew-KL and reverse-skew-KL implementations was the manual tuning of interpolation coefficients (alpha-1 for skew-KL and alpha-2 for skew-reverse-KL). DistiLLM-2 introduced an automatic updating rule for these coefficients, leveraging a first-order Mercator series approximation. This innovation eliminates manual tuning, allowing the balance between interpolation terms to adapt dynamically as training progresses and the teacher-student distribution gap narrows.
• Adaptive Skew-Reverse-KL Weighting: Recognizing the challenge of exactly matching teacher and student distributions on TGOs, especially with limited post-training data, DistiLLM-2 implemented a curriculum-based approach. It gradually increases the coefficient (beta-T) of the skew-reverse-KL term over training iterations. This modification strategically emphasizes fitting the student to the teacher's "tail" distribution using SGOs in later stages, effectively mitigating issues arising from inherent architectural and pre-training differences between teacher and student.

1. Superior Performance and Broad Applicability: Experimental evaluations demonstrated that DistiLLM-2 consistently achieves the highest inference speedup and better token-level alignment compared to other knowledge distillation baselines. This was observed across a diverse range of teacher-student pairs, including Qwen, Mistral, Gemma, and DeepSeek, and validated on multiple benchmarks such as instruction-following, math, and code generation. Moreover, the framework's versatility was showcased by its successful extension to preference alignment and vision-language models, and its compatibility with other efficiency techniques like network quantization and speculative decoding.

In essence, DistiLLM-2's key findings highlight the power of a synergized approach to knowledge distillation, where intelligent data selection meets sophisticated, adaptive loss function design. By making these relatively simple yet profound changes, DistiLLM-2 significantly enhances the effectiveness and efficiency of distilling knowledge into smaller models, paving the way for more widespread and practical LLM deployment.

Technical Deep Dive

▶ Watch: KALD framework: Integrating skew-KL, reverse-KL, and data strategies (6:00)

The technical foundation of DistiLLM-2 lies in its innovative contrastive approach for LLM distillation (KALD), which is then further optimized through adaptive data creation and a curriculum-based loss formulation. This section elaborates on these core technical components.

KALD: The Contrastive Foundation

The initial inspiration for KALD came from the success of contrastive methods like DPO and IPO in preference optimization. These methods operate by not only reinforcing desired outcomes but also actively suppressing undesired ones. The DistiLLM-2 team sought to apply this "contrastive" philosophy to knowledge distillation. However, directly porting DPO's loss function, as seen in DPKD, proved problematic for SLMs, often causing them to lose valuable pre-trained information.

This led to a re-evaluation of how contrastive learning could be implicitly achieved in KD. The solution proposed in KALD leverages the distinct behaviors of two fundamental information-theoretic measures: forward KL divergence (skew-KL) and reverse KL divergence (skew-reverse-KL).

• Skew-KL (Forward KL): This divergence, often denoted as $KL(P || Q)$, measures how much the distribution $Q$ differs from $P$. In the context of KD, if $P$ is the teacher's distribution and $Q$ is the student's, minimizing skew-KL encourages the student to place high probability mass where the teacher places high probability mass. It effectively "pulls up" the likelihood of the "head" part of the teacher's distribution. For KALD, the skew-KL component is used with Teacher-Generated Outputs (TGOs), which are considered the "desired" responses. It encourages the student model to match the teacher's high-probability predictions for these high-quality outputs.

• Skew-Reverse-KL (Reverse KL): Conversely, reverse KL divergence, $KL(Q || P)$, measures how much $P$ differs from $Q$. Minimizing reverse KL prevents $Q$ from assigning high probability to events where $P$ assigns low probability. In KD, it "pushes down" the likelihood of the "tail" part of the teacher's distribution, or more accurately, discourages the student distribution from diverging too far from the teacher's when the teacher assigns low probability. For KALD, the skew-reverse-KL component is used with Student-Generated Outputs (SGOs), which are considered "undesired" responses in the sense that they represent potential divergences from the teacher. This term discourages the student from producing responses that significantly deviate from the teacher's distribution, especially in areas where the teacher is less confident or would not produce such an output.

By combining these two divergences, KALD constructs a "response-aware loss formulation." The skew-KL term, applied to TGOs, ensures the student learns to produce high-quality outputs similar to the teacher. The skew-reverse-KL term, applied to SGOs, acts as a regularizer, preventing the student from generating low-quality or off-distribution responses. The talk highlights that this combined approach leads to faster convergence and better performance than using either divergence in isolation, effectively capturing the essence of a contrastive approach without the pitfalls of direct DPO application.

Data Creation Strategies for KALD

While KALD provides the loss function, the choice of data to feed these divergence terms is equally critical. The team extensively explored data creation strategies, initially experimenting with interpolations of TGO and SGO using speculative decoding. This involved tuning hyperparameters to balance "response quality" (how good the generated text is) and "on-policy behavior" (how much it reflects the student's current generation style).

However, their empirical findings converged on a surprisingly effective and straightforward strategy:

• For the skew-KL component, use Teacher-Generated Output (TGO) as the "desired" response (YT in their internal formulas). This ensures the student is pulled towards the high-quality outputs of the powerful teacher.
• For the skew-reverse-KL component, use Student-Generated Output (SGO) as the "undesired" response (YS in their internal formulas). This helps the student learn from its own mistakes or divergences from the teacher's ideal distribution.

Additionally, the research noted that YT (the high-quality desired response) could be further enhanced by using outputs generated from LLMs superior to even the designated teacher model, ensuring the student aims for the highest possible quality.

DistiLLM-2: Loss Formulation Refinements

DistiLLM-2 builds upon KALD by introducing two key refinements to the loss formulation, addressing practical challenges and further optimizing the distillation process:

1. Curriculum-Based Adaptive Interpolation Coefficients:

• Problem: In the previous iterations of skew-KL and reverse-skew-KL, the interpolation coefficients (referred to as alpha-1 for skew-KL and alpha-2 for skew-reverse-KL in the talk) had to be manually tuned. This manual tuning is cumbersome and suboptimal, as the optimal balance between interpolation terms likely changes throughout training. The observation was that in later training stages, as the student model improves and the gap between teacher and student distributions narrows, a smaller or similar alpha value might suffice.
• Solution: DistiLLM-2 introduces an automatic updating rule for alpha-1 and alpha-2. This rule is based on a first-order Mercator series approximation, allowing the coefficients to adapt dynamically during training. While the specific formula was presented visually in the talk, the core idea is to progressively adjust these interpolation terms, eliminating the need for manual tuning and ensuring an optimal balance as the student model evolves.

1. Adaptive Skew-Reverse-KL Weighting for Distribution Matching:

• Problem: Even with optimized loss functions and data, exactly matching the distributions of a powerful teacher and a smaller student on TGOs can be challenging. This difficulty is exacerbated when using small training datasets, particularly during post-training fine-tuning. The inherent differences in model size, number of parameters, and the quality or breadth of their pre-trained corpora contribute to this mismatch.
• Solution: To address this, DistiLLM-2 adopts a curriculum-based strategy that specifically focuses on the skew-reverse-KL term in later stages of training. The framework gradually increases the coefficient of skew-reverse-KL, denoted as beta-T for each training step T. This modification strategically emphasizes fitting the student's "tail" distribution using SGOs as training progresses. By giving more weight to preventing divergence from the teacher's low-probability areas, especially when the student is more mature, DistiLLM-2 helps to implicitly align the student's overall distribution more closely with the teacher's, compensating for fundamental differences that direct TGO matching might struggle with.

These refinements, though seemingly simple, are highly effective. They enhance the robustness, efficiency, and performance of the distillation process, culminating in the DistiLLM-2 framework that jointly optimizes data usage and loss dynamics for superior SLM training.

Experimental Setup & Results

▶ Watch: Adaptive alpha and curriculum-based beta for DistiLLM-2 (10:00)

The effectiveness and broad applicability of DistiLLM-2 were rigorously evaluated through a comprehensive experimental setup, covering various aspects of language model performance and efficiency.

Teacher-Student Pairs

The experiments utilized a diverse set of prominent LLMs as teacher models and corresponding SLMs as student models. The specific teacher-student pairs included:

• Qwen
• Mistral
• Gemma
• DeepSeek

This selection ensures that the findings are not specific to a single model architecture or family but are generalizable across different state-of-the-art LLMs, demonstrating the robustness of DistiLLM-2.

Evaluation Benchmarks

To assess the student models' capabilities, evaluations were conducted across multiple critical benchmarks, reflecting various aspects of language understanding and generation:

• Instruction-following ability: Measures how well the student model can interpret and execute complex instructions, a crucial capability for modern LLMs.
• Math generation: Evaluates the model's proficiency in solving mathematical problems and generating correct numerical or logical outputs.
• Code generation: Assesses the model's ability to generate syntactically correct and functionally sound code based on natural language prompts.

These benchmarks provide a holistic view of the distilled SLMs' performance, showcasing their enhanced reasoning and generative capabilities.

Broader Applicability and Efficiency Techniques

Beyond core language tasks, DistiLLM-2's versatility was explored in advanced domains and in conjunction with other efficiency techniques:

• Preference Alignment: The framework was extended to tasks involving preference alignment, demonstrating its potential in fine-tuning models to align with human preferences, similar to DPO but within a distillation context.
• Vision-Language Models (VLMs): The applicability of DistiLLM-2 was also demonstrated for vision-language models, suggesting its utility beyond purely text-based domains, indicating a powerful generalizable mechanism for knowledge transfer.
• Network Quantization: The paper, though not detailed in the talk, includes results on combining DistiLLM-2 with network quantization. This is a critical technique for reducing model size and accelerating inference by using lower-precision numerical representations.
• Speculative Decoding: The talk specifically highlighted the combination of DistiLLM-2 with speculative decoding. This inference-time optimization technique uses a small, fast "drafter" model to pre-generate tokens, which are then verified by a larger, more accurate "verifier" model. The research demonstrated that distillation from verifiers to drafters, improved by DistiLLM-2, significantly enhances speculative decoding efficiency. This implies that DistiLLM-2 can produce higher-quality drafter models, leading to greater inference speedups.

Headline Results

The experimental results consistently positioned DistiLLM-2 as a leading knowledge distillation method:

• Highest Inference Speedup: Compared to other KD baseline methods, DistiLLM-2 achieved the highest inference speedup. While specific numerical speedup factors were not detailed in the talk (referencing the poster/paper for these details), this indicates a substantial improvement in the efficiency of the distilled student models.
• Better Token-Level Alignment: The superior inference speedup is attributed to better token-level alignment achieved by DistiLLM-2. This implies that the student model's output distribution, particularly at the individual token level, more closely matches that of the teacher, leading to higher confidence in predictions and more efficient verification in techniques like speculative decoding.

In summary, the experimental evaluation firmly establishes DistiLLM-2 as a highly effective and versatile knowledge distillation framework. Its ability to train smaller models that not only perform well across diverse tasks but also demonstrate significant inference speedups, even when combined with other efficiency techniques, underscores its practical value for deploying high-performance AI in resource-constrained environments.

Practical Implications

▶ Watch: Experimental evaluation and significant effectiveness (12:00)

DistiLLM-2 offers profound practical implications across various stakeholders in the AI/ML ecosystem, from individual practitioners to large infrastructure teams. Its focus on creating effective and efficient small language models (SLMs) from powerful LLMs directly addresses critical deployment challenges.

For Practitioners and Model Builders

• Enabling On-Device and Edge AI: The most immediate implication is the ability to deploy highly capable SLMs in resource-constrained environments such as smartphones, IoT devices, and edge servers. Prior to such advanced distillation methods, achieving LLM-grade performance on these platforms was often infeasible due to memory, power, and computational limitations. DistiLLM-2 makes advanced language understanding and generation accessible where it was once out of reach.
• Reduced Inference Costs and Latency: For any application relying on LLMs, moving from large teacher models to distilled student models can dramatically reduce inference costs and latency. This translates to lower cloud computing bills, faster response times for users, and the ability to scale applications more economically. The demonstrated "highest inference speedup" directly contributes to this.
• Robust and Effective Distillation: DistiLLM-2 provides a robust and reliable framework for model builders. By jointly optimizing data creation and loss formulation, it offers a more principled and less heuristic approach to KD. The adaptive and curriculum-based strategies reduce the burden of manual hyperparameter tuning (e.g., alpha coefficients), accelerating the development cycle and making the distillation process more accessible even for those without deep expertise in KD nuances.
• Guidance on Data Strategy: The finding that Teacher-Generated Output (TGO) for skew-KL and Student-Generated Output (SGO) for skew-reverse-KL is most effective provides clear, actionable guidance for practitioners on how to prepare their distillation datasets. This simplifies a often complex aspect of KD.
• Broader Model Compatibility: The demonstrated applicability to vision-language models and preference alignment suggests that DistiLLM-2 is not just for text generation but can be a general-purpose tool for distilling knowledge across various AI modalities and objectives.

For Infrastructure and Deployment Teams

• Optimized Resource Utilization: Infra teams can achieve significantly better utilization of their computational resources. Deploying SLMs instead of LLMs means fewer GPUs or TPUs are needed for inference, or existing hardware can handle a much higher query volume.
• Synergy with Existing Efficiency Techniques: The compatibility with network quantization and speculative decoding is a major advantage. This means DistiLLM-2 can be integrated into existing MLOps pipelines that already employ these techniques, leading to cumulative efficiency gains. For instance, a DistiLLM-2 distilled model can serve as a highly effective drafter in speculative decoding, further boosting throughput.
• Scalability and Reliability: With smaller, faster models, deployments become more scalable and less prone to bottlenecks. This enhances the reliability and responsiveness of AI-powered services.

Tradeoffs and Limitations

While highly beneficial, DistiLLM-2, like any advanced technique, comes with certain tradeoffs and considerations:

• Increased Complexity in Loss Formulation: While alpha coefficients are automated, the overall loss function is more intricate than vanilla KD, involving multiple KL divergence terms and adaptive coefficients (beta-T). This requires a deeper understanding of the underlying principles for effective debugging or further customization.
• Reliance on Teacher Model Quality: The success of distillation is fundamentally tied to the quality of the teacher model. If the teacher is flawed or biased, these issues can be transferred to the student.
• Data Generation Requirements: While the data strategy is clear, generating high-quality TGOs and SGOs still requires computational resources and careful curation, especially for specific domains or tasks.
• Computational Cost of Distillation: The distillation process itself, particularly with sophisticated loss functions and potentially large datasets, will still incur computational costs for training, though this is a one-time cost for significant long-term inference savings.
• Generalization Beyond Benchmarks: While tested on diverse benchmarks, real-world deployment always presents unique challenges. Practitioners will need to validate performance on their specific production data and tasks.
• Inherent Teacher-Student Gap: Despite the beta-T curriculum, the fundamental differences in parameter count and pre-training data between a large teacher and a small student cannot be entirely eliminated. There will always be a performance gap, albeit significantly narrowed by DistiLLM-2.

In conclusion, DistiLLM-2 offers a powerful and practical solution for democratizing access to LLM capabilities. By enabling efficient deployment of SLMs without significant performance compromise, it paves the way for a new generation of AI applications that are faster, cheaper, and more ubiquitous.

Key Takeaways

• DistiLLM-2 is a novel knowledge distillation (KD) framework that introduces a contrastive approach to effectively and efficiently transfer knowledge from large language models (LLMs) to small language models (SLMs).
• It uniquely optimizes both loss formulation and data creation simultaneously, a synergy overlooked by prior KD research, leading to superior student model performance.
• The core method, KALD, leverages forward KL (skew-KL) and reverse KL (skew-reverse-KL) divergences to implicitly achieve contrastive learning: skew-KL pulls the student towards high-probability teacher outputs (TGOs), while skew-reverse-KL pushes down student divergence from low-probability teacher outputs (SGOs).
• DistiLLM-2 refines KALD with adaptive, curriculum-based loss components: It automatically adjusts interpolation coefficients (alpha-1, alpha-2) using a first-order Mercator series approximation and gradually increases the skew-reverse-KL coefficient (beta-T) during training to better handle inherent teacher-student distribution gaps.
• Experimental results demonstrate significant performance gains: DistiLLM-2 achieves the highest inference speedup and better token-level alignment compared to other KD baselines, tested across diverse teacher-student pairs (Qwen, Mistral, Gemma, DeepSeek) and benchmarks (instruction-following, math, code generation).
• The framework exhibits broad applicability and compatibility: It extends to preference alignment and vision-language models, and integrates effectively with other efficiency techniques like network quantization and speculative decoding, making it a versatile tool for various AI deployment scenarios.

About the Speaker(s)

The talk was presented by Jongwoo Ko, who is affiliated with KAIST AI. This research is a collaborative effort between KAIST AI and Microsoft. The co-authors listed are Tianyi Chen, Sungnyun Kim, Tianyu Ding, Luming Liang, Ilya Zharkov, and Se-Young Yun, indicating a strong inter-institutional team contributing to this advanced work in large language model distillation.