PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jongwoo Ko, Tianyi Chen, Sungnyun Kim, Tianyu Ding, Luming Liang, Ilya Zharkov, Se-Young Yun

arXiv:2503.07067v2cs.CLcs.AIcs.LG

TL;DR

LLM distillation seeks to compress resource-intensive models, but applying the same loss treatment to teacher- and student-generated data overlooks their differing roles. DistiLLM-2 combines contrastive loss dynamics with curated data and curriculum-based learning, achieving state-of-the-art performance across several text-generation tasks and extending to preference alignment and vision-language settings. Its contrastive design addresses reward-hacking behavior observed when DPO is directly applied to distillation.

  • Problem

    LLMs are costly to deploy, while prior distillation strategies overlook the synergy between loss formulations and teacher- versus student-generated data.

  • Method

    DistiLLM-2 combines contrastive loss, curated data, and curriculum-based learning, using distinct loss dynamics for teacher- and student-generated responses.

  • Results

    DistiLLM-2 achieves state-of-the-art performance across instruction-following, mathematical reasoning, and code-generation tasks, with applications in preference alignment and vision-language models.

  • Takeaways & Limitations

    The framework provides a contrastive approach for aligning teacher and student models across varied data types and improving compressed-model effectiveness.

  • Takeaways & Limitations

    Directly applying DPO to distillation can cause reward hacking and excessively reduce the likelihood of student responses, leading to degenerate sentences and loss of pretrained information.

Abstract

from arXiv · show

Despite the success of distillation in large language models (LLMs), most prior work applies identical loss functions to both teacher- and student-generated data. These strategies overlook the synergy between loss formulations and data types, leading to a suboptimal performance boost in student models. To address this, we propose DistiLLM-2, a contrastive approach that simultaneously increases the likelihood of teacher responses and decreases that of student responses by harnessing this synergy. Our extensive experiments show that DistiLLM-2 not only builds high-performing student models across a wide range of tasks, including instruction-following and code generation, but also supports diverse applications, such as preference alignment and vision-language extensions. These findings highlight the potential of a contrastive approach to enhance the efficacy of LLM distillation by effectively aligning teacher and student models across varied data types.

1. Introduction

LLM distillation is motivated by the need to reduce inference costs while preserving model performance. DistiLLM-2 addresses limitations in prior loss and data strategies with contrastive, asymmetric training and additional curation and curriculum mechanisms.

  • LLMs require extensive inference resources, motivating parameter reduction that preserves performance for practical deployment.
  • Prior distillation work largely separates innovations in loss design from training-data curation, despite limitations in standard KL divergence.
  • DistiLLM-2 introduces CALD, which applies distinct loss functions to teacher- and student-generated samples to exploit their synergy with data types.
  • The framework extends CALD with optimized dataset curation and curriculum-based adaptive loss mechanisms intended to guide practitioners.
  • DistiLLM-2 achieves state-of-the-art performance across instruction-following, mathematical reasoning, and code-generation tasks, while supporting preference alignment and vision-language extensions.

2. Backgrounds

Knowledge distillation compresses large models into smaller students, but LLM distillation depends on both loss design and training-data curation. Prior work explored KL variants and adaptive off-policy data reuse, while this work further refines the objectives with a contrastive approach and batch data collection.

  • Distillation for LLMs: Knowledge distillation enables smaller student models to match larger teachers and has been adapted to improve LLM scalability.Prior work also used student-generated outputs as distillation data and explored on-policy objectives such as reverse KL.
  • Loss functions: KL distillation minimizes divergence between teacher and student response distributions and can be decomposed into token-wise losses.The teacher and student distributions are denoted p(y|x) and qθ(y|x), respectively.
  • Loss functions: Forward and reverse KL can suffer from mode-averaging and mode-collapsing, motivating skew KL and skew reverse KL objectives.These alternatives were proposed to address limitations of the conventional KL variants.
  • Loss functions: Skew KL achieved faster convergence and better performance than baselines including MiniLLM and GKD, with empirical and theoretical support.The passage groups skew KL with skew reverse KL but specifically reports the performance comparison for SKL.
  • Data curation: Adaptive off-policy distillation reuses student-generated outputs through a replay buffer to balance computational efficiency and training efficacy.The approach addresses training inefficiency and low-quality student-generated outputs in on-policy setups.
  • Connection to this work: This work combines contrastive objective refinement with batch collection of student-generated outputs before each training epoch.The batch setup is compatible with vLLM, and preliminary results report greatly reduced sample-generation cost with minimal impact on student performance.

3. Method: DISTILLM-2

DISTILLM-2 extends knowledge distillation with a contrastive loss that assigns different objectives to teacher- and student-generated responses, supplemented by data curation and adaptive curricula. Its components address reward hacking, training–inference mismatch, and response-type synergy while improving convergence and student performance.

  • Contrastive loss: DISTILLM-2 combines SKL for teacher responses with SRKL for student responses, balancing them through coefficient β in [0, 1].The loss is designed to increase alignment with teacher-generated responses while regularizing student-generated responses.
  • Motivation: DPKD can reward-hack by excessively decreasing the student likelihood of student responses, producing degenerate sentences and discarding pretrained information.The reported NLL values are 91.25 for student responses and 20.29 for teacher responses.
  • Contrastive loss: Using different loss types for different response types converges faster and performs better than applying SKL or SRKL exclusively to the same response type.Figure 1(c) reports greater effectiveness and higher ROUGE-L for CALD with SKL and SRKL.
  • Contrastive loss: CALD uses a linear, token-decomposable formulation with explicit probability weighting, which regularizes excessive decreases in student-response likelihood.This distinguishes CALD from the log-sigmoid formulation used in DPKD and DPO.
  • Data curation: Data curation favors strong teacher guidance for SKL and low-quality student-generated responses for SRKL, rather than uniformly minimizing training–inference mismatch.Replacing teacher responses with speculative-decoding responses performs best at ε = 1.0 for SKL, whereas SRKL performs best at ε = 0.0.
  • Adaptive learning: DISTILLM-2 adds curriculum-based α updates and gradually increases β, improving performance by reducing difficult teacher imitation and emphasizing feedback from student-generated outputs.Moderate α values such as 0.1–0.3 are generally robust, but optimal values vary across teacher–student pairs and epochs.

4. Experiments

DISTILLM-2 is evaluated across instruction-following, mathematical reasoning, and code generation, where it generally outperforms competing distillation methods. The experiments use benchmark evaluations and compare the method with established baselines.

  • General Instruction-Following: Instruction-following evaluation compares DISTILLM-2 with other baselines on AlpacaEval, Evol-Instruct, and UltraFeedback using LLM-based judging.The setup uses 50k prompts sampled from UltraChat200k and teacher- and student-generated responses.
  • General Instruction-Following: DISTILLM-2 performs best in most instruction-following evaluation setups, except for Danube2-1.8B on Evol-Instruct.
  • General Instruction-Following: +2.34%, +1.95%, and +4.53% average gains over the second-best methods are reported for Qwen2-1.5B, Danube2-1.8B, and Gemma2-2B, respectively.The results are presented as average improvements across the evaluated instruction-following benchmarks.
  • Mathematical Reasoning: On GSM8K and MATH, DISTILLM-2 achieves higher performance than other baselines in both Qwen2 and Qwen2.5 experimental setups.The Qwen2.5 student also outperforms the Qwen2 teacher on MATH.
  • Code Generation: Code-generation experiments evaluate DISTILLM-2 against GKD and DistiLLM on HumanEval and MBPP.The study uses Qwen2.5-Coder and DeepSeek-Coder teacher models with smaller student models.
  • Code Generation: Across HumanEval and MBPP, DISTILLM-2 consistently outperforms GKD and DistiLLM.The reported comparison attributes the outcome to combining SKL or SRKL with harmonized response types.

5. Additional Ablation Study

The ablation study examines DISTILLM-2 components, training-data size, and teacher capacity. Results indicate incremental gains from the proposed components and consistent advantages across training sizes and teacher scales.

  • Component Analysis: The component analysis tests contrastive learning, a larger β parameter, and curriculum-based updates to α.These components are progressively added to DistiLLM for comparison.
  • Component Analysis: Each added component produces incremental performance gains, bringing DistiLLM closer to DISTILLM-2’s overall effectiveness.
  • Training Size: Across instruction-following, mathematical reasoning, and code-generation tasks, DISTILLM-2 consistently outperforms GKD and DistiLLM at different training-data sizes.Figure 3 compares performance as the amount of training data varies.
  • Capacity Gap: The capacity-gap analysis varies Qwen1.5 teacher sizes from 1.8B to 14B parameters while using a 0.5B student.The capacity gap is described as the challenge created by the size difference between teacher and student models.
  • Capacity Gap: As teacher size increases, DISTILLM-2 shows monotonic improvement and consistently outperforms other baselines.The study reports stronger handling of capacity gaps than the previous version, particularly with 7B and 14B teachers.

6. Broader Impacts

The paper extends DISTILLM-2 beyond standard language-model distillation to preference alignment, quantized-model recovery, speculative decoding, and vision-language models. These applications report improved alignment, restored performance, or faster inference relative to relevant baselines.

  • Applications: DISTILLM-2 is presented for preference alignment, quantized-model recovery, speculative decoding, and vision-language distillation.The paper describes these as diverse applications of the proposed knowledge-distillation approach.
  • Preference Alignment: Replacing standard SFT with language-model KD is evaluated as the first stage before DPO-based preference fine-tuning.The comparison examines how the resulting policy models perform after the preference-optimization stage.
  • Preference Alignment: DISTILLM-2 achieves substantially higher WR in preference fine-tuning, more than doubling WR for Qwen2-1.5B and showing a similar improvement for Gemma2-2B.The results support using DISTILLM-2 to build reference models for subsequent preference alignment.
  • Vision-Language Models: DISTILLM-2 is applied to vision-language distillation with LLaVA-1.5-7B as teacher and TinyLLaVA-1.4B as student, evaluated on OK-VQA and TextVQA.The training data comes from the 83K-prompt RLAIF-V-Dataset.
  • Quantized Models: KD methods significantly improve the performance of 4-bit quantized models when LoRA replaces regular SFT.The paper evaluates DISTILLM-2 for restoring performance after post-training quantization.
  • Speculative Decoding: DISTILLM-2-trained drafters achieve higher speculative-decoding speedups than SFT- and DistiLLM-trained drafters with Phi3.5-mini and Phi3-medium verifiers.The comparison is summarized for draft models evaluated with both verifier models.

7. Conclusion

The paper presents DISTILLM-2 as a distillation framework combining contrastive loss, curated data, and curriculum-based learning. Across language and vision-language settings, it reports stronger performance, sample efficiency, alignment, and generalization.

  • Conclusion: DISTILLM-2 combines contrastive loss, curated data, and curriculum-based learning in a distillation framework for large language models.The conclusion describes the framework as differentiating teacher and student outputs.
  • Conclusion: Experiments across instruction following, mathematical reasoning, and code generation report state-of-the-art performance with improved sample efficiency.The conclusion also states reduced reliance on expensive preference-labeled data.
  • Conclusion: The paper reports enhanced reasoning ability and robustness in language and vision-language models through high-quality explanations and contrastive objectives.

Impact Statement

The paper describes applications of DISTILLM-2 beyond standard distillation, including preference alignment, multimodal models, recovery of compressed models, and faster speculative decoding. It also identifies broader-domain use as a direction for future work.

  • Impact Statement: DISTILLM-2 can replace vanilla SFT before PPO or preference-optimization methods such as DPO, IPO, and SimPO.The authors suggest this could support preference alignment with stronger reference models.
  • Impact Statement: DISTILLM-2 demonstrates potential adaptability to vision-language models because many multimodal LLMs are built on LLMs.
  • Impact Statement: DISTILLM-2 can substantially improve quantized LLMs, making them competitive with their uncompressed counterparts.The authors also suggest applicability to other compression techniques.
  • Impact Statement: DISTILLM-2 can enhance speculative decoding by improving alignment between draft and target models.The authors connect this application to systems that use diverse models within one framework.
  • Impact Statement: The authors believe DISTILLM-2 could be used across a broader range of fields than those demonstrated.They mention integrating reasoning abilities such as those showcased by DeepSeek-R1 as future work.

Supplementary Material

Prior LLM distillation research has explored model- and data-centered approaches, while contrastive methods have mainly been developed for preference alignment and reasoning. Extending contrastive schemes to LLM distillation remains relatively underexplored.

  • Supplementary Material: LLM distillation research has focused largely on designing new losses or curating training data.
  • Supplementary Material: Contrastive approaches such as DPO use different learning strategies for two distinct responses in preference alignment or reasoning.
  • Supplementary Material: Few works had extended contrastive approaches to knowledge distillation for LLMs.

B. Derivation for Mathematical Analysis

This supplementary material explains how KL-family divergences shape student probabilities and motivates the loss design used in DISTILLM-2. It also describes a first-order approximation, implementation stabilization, and the datasets and evaluation tasks used experimentally.

  • B. Derivation for Mathematical Analysis: The f-divergence framework defines KL and reverse KL through different generating functions, while skew variants use interpolated forms.KL uses f(t) = t log t, reverse KL uses f(t) = −log t, and α-skew reverse KL uses f(α)(t) = −log((1−α)t + α).
  • B. Derivation for Mathematical Analysis: Forward-KL-like terms pull student probabilities upward where the teacher assigns substantial probability.The same pulling-up property is attributed to SKL.
  • B. Derivation for Mathematical Analysis: Reverse-KL-like terms push student probabilities downward where the teacher assigns low probability.The same pushing-down property is attributed to SRKL, including for student responses in the distribution tail.
  • B. Derivation for Mathematical Analysis: The paper claims its mathematical explanation applies to a more general setup than a prior unimodal-Gaussian analysis and covers both skew and ordinary KL variants.
  • B. Derivation for Mathematical Analysis: A first-order Mercator approximation expresses the skew divergences using α, p(y|x), and qθ(y|x), enabling sample-wise selection of α.The implementation uses mini-batch allocation and clipping, with bounds 0.1 and 0.01, to improve stability.
  • C.1. Dataset Description: Experiments cover instruction-following, mathematical reasoning, and code-generation datasets, including UltraChat200k, AlpacaEval, Evol-Instruct, UltraFeedback, MetaMathQA, GSM8K, MATH, WizardCoder, HumanEval, and MBPP.The listed datasets span training and evaluation settings across these task categories.

C.3. Evaluation

The evaluation samples model outputs under task-specific decoding settings and compares instruction-following responses with LLM-as-a-Judge prompts. Mathematical reasoning and code generation use greedy sampling, while instruction-following uses temperature-based sampling.

  • C.3. Evaluation: Instruction-following responses are sampled on one NVIDIA A100 80GB GPU with temperature 0.8, top-p 0.95, and a maximum length of 512.
  • C.3. Evaluation: LLM-as-a-Judge evaluation uses a pairwise comparison prompt with temperature 0.7.For AlpacaEval, responses are compared against officially released text-davinci-003 outputs.
  • C.3. Evaluation: Mathematical-reasoning and code-generation evaluations use greedy sampling with a maximum length of 1024.Code generation is evaluated with the EvalPlus framework.
  • C.3. Evaluation: The judge is instructed to compare two assistants on helpfulness, relevance, accuracy, depth, creativity, and detail, then output A, B, or C.

D.1. Comparison on On-policy Setup

The section evaluates DistiLLM methods across on-policy configurations, distillation baselines, and an instruction-tuned student setting. Results indicate that the batched setup preserves performance efficiently, while DISTILLM-2 performs strongly with limited training iterations.

  • Experimental setup: The experiments compare recent on-policy distillation baselines, adaptive on-policy configurations, DPKD, and DPO under the DistiLLM codebase.The setup follows Ko et al. (2024) for GPT-2 trained on databricks-dolly-15k.
  • On-policy comparison: The batched on-policy setup avoids severe performance degradation despite its significant efficiency compared with fully on-policy and off-policy setups.Students are evaluated on databricks-dolly-15k using ROUGE-L.
  • Baseline comparison: DPKD performs much worse than its reported values and is excluded from the main baselines because it is prone to reward hacking.The authors still provide results for DPO except in the omitted portion of the passage.
  • Instruction-tuned student: The instruction-tuned student experiment uses Gemma-2-2B-it with Gemma-2-9B-SimPO as teacher and trains for one epoch of 200 iterations.This setup otherwise matches the base-model experiment.
  • Results: DISTILLM-2 achieves higher UltraFeedback performance than the other student models and even exceeds the teacher there, despite very limited training iterations.The authors attribute these results to fast convergence of the contrastive approach.
Loading 2503.07067v2…