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

We thank the reviewer for their constructive feedback. Below we address each concern in detail. Look forward to any further discussions.

Q1. Lines 152–156 might be confusing – rephrasing suggested

A1. Thank you for the suggestion. We will revise Lines 152–156 to improve clarity. This paragraph aims to illustrate the relation between contrastive learning in preference alignment and our approach to knowledge distillation: increasing the student model's likelihood of generating outputs aligned with teacher responses, and decreasing it for those resembling weaker student responses, using distinct loss functions.

To make this point clear, we plan to revise the paragraph as follows:

“Similarly, we incorporate this contrastive concept from preference optimization into KD by assigning different loss functions to different types of responses: encouraging the student model to assign higher likelihood to high-quality responses generated by the teacher ($y_t$) -- $q_\theta(y_t|x)$ -- while reducing the likelihood of lower-quality student responses ($y_s$) -- $q_\theta(y_s|x)$ -- that deviates from the teacher.”

We appreciate your constructive comments. We have rephrased the question for ease of reference and provide our corresponding responses below. Look forward to any further discussions.

Q1. Inclusion of acceptance rates in speculative decoding analysis for completeness

A1. Thank you for the suggestion. To provide a more concrete answer, we evaluated the acceptance rates using the vLLM [1] framework, which supports detailed speculative decoding metrics. The results of acceptance rates are presented as follows:

We could see that DistiLLM-2 significantly boosted the acceptance rates leading to a higher speed-up via speculative decoding.

[1] Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP. 2023

We thank the reviewer for their constructive feedback. Below we address each concern in detail. Look forward to any further discussions.

Q1. Empirical validation of the approximation in Equation (7) for sequence-level probabilities

A1. Thank you for pointing out the potential mismatch between our first-order Mercator approximation in Equation (7) and actual sequence-level probabilities $p(y|x)$. Indeed, this approximation assumes $p(y|x) \simeq 1$, which can be problematic for long sequences where $p(y|x)$ might vanish.

In our implementation, we compute probabilities at the token level, clip extreme values to avoid outliers, and average them over the sequence, yielding a reasonably scaled alpha that is shared across all tokens.

The corresponding code (in src/distillm_trainer.py, lines 1161–1164) is:

anchor = (1 - base_alpha_1) * logp_logq
logps_logqs = (
    (tea_per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
).exp() - (
    (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
).exp()  # sentence-level
alpha_1 = torch.clip(
    1 - anchor / (logps_logqs + 1e-5),
    min=1e-2, max=base_alpha_1
).unsqueeze(-1).unsqueeze(-1)

This design makes the approximation invariant to sequence length, effectively resolving the concern about very small $p(y|x)$ values in long sequences. Notably, we want to highlight that this first-order approximation has advantage in closed-form $\alpha$ updates for efficient, per-sample adaptation.

Additionally, we provide empirical results of token-level probabilities on teacher ($y_t$) and student ($y_s$) responses to quantify the first-order approximation using the teacher (Mistral-7B) and student (Danube2-1.8B) models. To better characterize the distribution of token-level probabilities, we report the first (Q1), second (Q2, median), and third (Q3) quartiles for both models:

The results show that most of the token-level probability values are close to 1, which ensures small approximation error of our first-order approximation. We appreciate the reviewer’s insightful comment and will update the paper accordingly to better explain this implementation and its practical implications.

Q2. CALD term appears without explanation – unclear terminology

A2. Thanks for the question. CALD stands for Contrastive Approach for LLM Distillation, which is first introduced in lines 60-61 in our main manuscript. To improve the readability and clarity, we will also define the term at the beginning of Section 3.1.2. We will additionally refine Section 3 and add more descriptive languages right after the definition (Equation (5)).

Thank you for your insightful feedback. We have rephrased your comments for simpler reference and have included our respective responses. Look forward to any further discussions.

Q1. Use of models up to 9B; interest in evaluating scalability to larger models

Thank you for the insightful comment. In the evaluation, we already demonstrated the effectiveness of DistiLLM-2 over varying sizes of teacher and student models. For teacher models, we have included larger backbones with over 9B parameters, such as Qwen1.5-14B in Appendix D.2 (Table 10) and Phi-3-Medium-14B in Appendix D.4 (Table 12).

For larger student models, we are working on conducting more such experiments. Although we might not be able to report the results in a timely manner during rebuttal due to the resource cost, we will report them in the future.

Mathematically, as explained in Remark 1, DistiLLM-2 proposes a dedicated CALD objective function that could exhibit behavior similar to DPO in preference alignment. Given DPO’s wide validation across diverse architectures and the supporting results from our experiments, we are confident in the generality of the proposed approach.