DistiLLM: Towards Streamlined Distillation for Large Language Models

• Autoregressive KD must compress high-capacity language models while preserving performance and reducing inference cost or memory footprint.Source: §1 · p. 1
• Fixed-dataset KLD can cause mode averaging or mode collapse and training-inference distribution mismatch in autoregressive language models.Source: §1 · p. 1
• Existing student-generated-output methods lack standardized objectives and incur inefficiency through continual sequence generation and task-dependent loss selection.Source: §1 · pp. 1–2
• DistiLLM introduces skew KLD with theoretical and empirical support for stable gradients, lower approximation errors, faster convergence, and improved performance.Source: §1 · p. 2
• Its adaptive off-policy approach improves the efficiency of using student-generated outputs by adaptively managing their contribution during distillation.Source: §1 · p. 2
• DistiLLM achieves state-of-the-art student-language-model performance across generative tasks while delivering a 2.5 ∼4.3× training speedup over recent KD techniques.Source: §1 · p. 2

2.1. KD for Auto-regressive Generative LMs

• KD trains a student qθ(y|x) to match a fixed teacher p(y|x), using data from ground-truth or teacher-generated sequences.Source: §2.1 · p. 2
• KLD is widely used because it is simple and tractable, and sequence-level KLD can decompose into token-wise distillation.Source: §2.1 · p. 2
• The practical formulation approximates distribution matching under an assumption about the teacher distribution, while the exact KLD supports theoretical analysis.Source: §2.1 · p. 2

2.2. Pitfalls of Existing Distillation

• Asymmetric KLD can force a limited-capacity student to cover the teacher’s support, producing an overly smooth distribution through mode averaging.Source: §2.2 · p. 2
• Reverse KLD and generalized JSD have shown empirical success, but autoregressive KD still lacks a theoretically and experimentally grounded standard objective, causing task-dependent variability.Source: §2.2 · pp. 2–3
• Student-generated outputs reduce training-inference mismatch, but unfamiliar or inaccurate sequences can misguide the teacher and generate misleading feedback.Source: §2.2 · p. 3, Figure 1 · p. 3
• DISTILLM’s training pipeline combines fixed-dataset and replay-buffer samples, updates with skew or skew-reverse KLD, and adaptively adjusts the student-generated-output ratio using validation loss.Source: §2.2 · p. 3
• Generating student outputs at every iteration is computationally inefficient, with generation accounting for up to 80% of total training time.Source: §2.2 · p. 3

3. DISTILLM

• DISTILLM combines skew KLD with an adaptive off-policy approach to address objective instability, noisy student-generated feedback, training-inference mismatch, and SGO generation cost.The framework presents skew KLD and an adaptive SGO scheduler plus replay-based off-policy training as its two main components.Source: §3 · p. 4, §3.2 · p. 5

3.1. Skew (Reverse) KLD

• Skew KLD interpolates teacher and student distributions, reducing gradient norms by preventing the interpolation denominator from reaching zero and thereby producing more stable optimization.The same gradient analysis reports analogous trends for skew reverse KLD.Source: §3.1 · p. 4
• Skew KLD has a bounded mini-batch estimator norm, reducing approximation error between empirical and true divergence and supporting generalizability.Theoretical analysis links larger α to a lower L2 norm, while the overall α choice balances gradient and estimator scales.Source: §3.1 · pp. 4–5
• With α = 0.1, α-SKL and α-SRKL outperform KLD, RKLD, and JSD in the reported experiments.The paper attributes the advantage to more stable gradients and smaller approximation error, while noting that skew choices involve a trade-off.Source: §3.1 · pp. 4–5

3.2. Adaptive Off-policy Approach

• The adaptive SGO scheduler starts with low SGO probability and increases it using validation-loss changes, limiting noisy feedback while reducing training-inference mismatch.SGOs are selected with probability ϕ, whereas fixed-dataset samples are used with probability 1−ϕ.Source: §3.2 · p. 5
• The off-policy strategy stores student-generated outputs in a replay buffer and reuses them, improving SGO sample efficiency and saving resources relative to continually generating new data.The method replaces older buffer entries when capacity is reached and randomly samples from the stored pool.Source: §3.2 · pp. 5–6
• The replay ratio decreases over training: early stages use current SGOs with more replay, whereas later stages reuse stored SGOs with less replay to balance bias and efficiency.The design maintains a small SGO generation frequency λR by varying replay according to training progress.Source: §3.2 · p. 6
• The adaptive off-policy approach works especially well with skew losses, preserving efficacy after switching from on-policy to off-policy and improving training efficiency with negligible performance drop.The paper links this synergy to the rapid early convergence of SKL and SRKL; alternatives show performance drops under the same switch.Source: §3.2 · p. 6

4. Experiments

• DISTILLM is evaluated with SRKL and adaptive off-policy training against SFT, KD, SeqKD, ImitKD, MiniLLM, and GKD.The instruction-following evaluation uses five benchmarks and ROUGE-L plus GPT-4 feedback; the task-specific experiments use SAMSum and IWSLT 2017.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 7

4.1. Task-Agnostic Instruction-Following

• The experiments report ROUGE-L validation scores and adaptive off-policy comparisons using figures and tables that summarize loss-function and training-efficiency effects.The supplied visuals cover loss-function ROUGE-L, relative generation time, and task-specific ROUGE-L comparisons.Source: Table 1 · p. 7, Figure 6 · p. 7, Table 2 · p. 7
• DISTILLM outperforms state-of-the-art methods across diverse instruction-following teacher-student combinations and evaluation metrics.GKD and MiniLLM are less effective or can reduce performance for smaller students, especially on GPT-4 feedback and some ROUGE-L scores.Source: §4.1 · p. 7
• DISTILLM surpasses other baselines for OpenLLaMA2-7B to OpenLLaMA2-3B distillation and requires only 1.6× naive KD training time.Other methods require 3–7× naive KD training time in this larger-model setting.Source: §4.1 · p. 7

4.2. Text Summarization and Machine Translation

• DISTILLM outperforms other baselines on SAMSum and IWSLT 2017, although its margin is smaller in single-task scenarios than in general-purpose instruction-following.ImitKD exceeds GKD on SAMSum, whereas GKD exceeds ImitKD on IWSLT, indicating task-dependent baseline behavior.Source: §4.2 · pp. 7–8

5.1. Effect of Skew Divergence

• SKL and SRKL generally outperform KLD, RKLD, and JSD while achieving high validation ROUGE-L throughout training.Their sustained validation performance indicates rapid convergence and strong generalization in the reported GPT-2 XL to GPT-2 experiments.Source: §5.1 · p. 8

5.2. Effect of Adaptive Off-policy Approach

• The adaptive off-policy strategy improves performance across datasets with minimal performance loss and achieves 2.2× to 3.4× faster training than on-policy or mixed strategies.It balances noisy feedback against training-inference mismatch by selectively using student-generated outputs.Source: §5.2 · p. 8

5.3. Additional Ablation Studies on DISTILLM

• SKL and SRKL perform best at α = 0.1, while SRKL degrades more sharply than SKL as α increases above 0.1.The result agrees with the reported normalized L2-norm analysis.Source: §5.3 · pp. 8–9
• Combining the proposed off-policy approach with the proposed objectives is more effective than applying off-policy training to GKD or ImitKD.The reported comparison attributes the difference to the slower training speeds of JSD and KLD relative to SRKL.Source: §5.3 · p. 8, Table 4 · p. 8
• The adaptive SGO scheduler usually exceeds manually selected probabilities, with learned probabilities within 0.2 of the manually tuned optima.The table caption reports terminal probabilities of 0.4 for SKL and 0.3 for SRKL.Source: §5.3 · p. 8, Table 5 · p. 8
• DISTILLM remains relatively robust when distilled directly from pretrained student parameters without prior fine-tuning.This setting produces only a minor performance reduction compared with fine-tuned initialization, unlike other methods in the reported comparison.Source: §5.3 · p. 9, Table 6 · p. 9

6. Related Work

• Knowledge distillation transfers large teacher-model capabilities to smaller autoregressive language models to make compute-intensive applications more viable.Related approaches also use LLM-generated annotations or reasoning data to fine-tune smaller models.Source: §6 · p. 9
• Divergence-based methods match student and teacher generation distributions, while recent work explores alternative objectives and student-generated outputs.ImitKD uses student-generated outputs; subsequent methods introduce on-policy training, diverse divergences, and policy-gradient optimization.Source: §6 · p. 9
• DISTILLM extends this research with an effective and efficient framework that analyzes both the objective function and data utilization.Its stated contribution is to combine skew divergence with adaptive use of training data.Source: §6 · p. 9

7. Conclusion

• DISTILLM combines theoretically analyzed SKL with an adaptive off-policy approach that reduces noisy SGO feedback and improves sample efficiency using a replay buffer.The two components are presented as the framework’s central response to existing KD challenges.Source: §7 · p. 9
• Extensive experiments across generation tasks demonstrate superior performance alongside significant training-efficiency improvements.The conclusion summarizes the reported empirical outcome without specifying a single dataset or metric.Source: §7 · p. 9

• DISTILLM improves distillation performance and computational efficiency for autoregressive LMs, reducing reliance on high-end resources for training small LLMs.Source: Impact Statement · p. 9

A. Limitation

• The method focuses on (R)KLD, leaving potential gains from combining SKL with TVD or EMD as future work.The authors expect SKL’s rapid early convergence could complement the later-training performance of EMD or TVD.Source: A. Limitation · p. 13
• DISTILLM is primarily designed for supervised fine-tuning, so extending it to human-preference optimization remains future work.Source: A. Limitation · p. 13
• The method assumes teacher and student models share the same tokenizer, leaving cross-tokenizer knowledge transfer for future exploration.Source: A. Limitation · p. 13

B. Further Discussion on Skew KLD

• The supplementary discussion derives theoretical results for SKL and adds empirical analysis of its gradient behavior.Source: B. Further Discussion on Skew KLD · p. 13

B.1. Details for Gradient Analysis

• The derivation computes sample-wise gradients for KLD, SKL, RKLD, and SRKL to analyze optimization stability.Source: §B.1 · p. 13
• SKL consistently reduces gradient-coefficient variance as α increases, whereas SRKL has its smallest variance at α=0.1 and grows thereafter.The analysis links these patterns to SKL’s greater robustness across α values and to smaller gradient scales as α increases.Source: §B.1 · p. 14, Figure 9 · p. 14

B.2. Proof of Theorem 1

• The proof assumes regularity conditions under which α-SKL admits a variance bound for α<1/8, unlike unbounded KLD and RKLD.Source: §B.2 · p. 14
• The theoretical derivation combines f-divergence concentration bounds with bias-variance decomposition to obtain the stated estimator guarantees.Source: §B.2 · pp. 14–15

C. Details of DISTILLM Algorithm

• The adaptive scheduler increases the probability of using student-generated outputs when validation loss rises, using a tolerance ε and staged updates.The implementation sets ε=0.1, initial SGO probability 0, and N_ϕ=10.Source: C. Details of DISTILLM Algorithm · pp. 15–16
• Validation loss is used to trigger additional SGOs because increasing loss suggests overfitting and a training-inference mismatch.The scheduler is intended to improve versatility while limiting noisy feedback.Source: C. Details of DISTILLM Algorithm · p. 16, Figure 10 · p. 16

D. Experimental Setup

• Experiments cover instruction-following, text summarization, and machine translation tasks.The datasets include instruction-following benchmarks, CNN/DM, SAMSum, XSum, and IWSLT 2017 English-to-German translation.Source: §D.1 · p. 16

D.1. Dataset Description

• The dataset description includes diverse instruction-following benchmarks spanning multiple task types and evaluation settings.Source: §D.1 · p. 16

D.2. Training Details

• Instruction-following models use task-dependent learning-rate, batch-size, and epoch searches, with checkpoints selected by validation ROUGE-L.Source: §D.2 · p. 17
• DISTILLM remains effective without fine-tuned student initialization required by several SGO-based baselines.Source: §D.2 · p. 17
• Task-specific training uses ten epochs for summarization, two for IWSLT 2017, and reduced samples for XSum and CNN/DM because SGO methods increase training time.Source: §D.2 · p. 17

D.3. Evaluation

• Instruction-following evaluation samples five random seeds with temperature 1.0, while GPT-4 feedback uses the Zheng et al. (2023) prompt at temperature 0.7.Source: §D.3 · p. 17, Figure 11 · p. 17, Figure 12 · p. 18
• Task-specific evaluation uses greedy decoding with maximum length 128, scoring summarization by ROUGE-L and translation by BLEU.Source: §D.3 · p. 17

E.1. Full Results of Instruction-Following (Fig. 5)

• DISTILLM outperforms MiniLLM and GKD on most instruction-following datasets and metrics across model sizes, achieving the best performance in most cases.Source: §E.1 · pp. 17–18

E. Additional Results

• The supplementary results report ROUGE-L and BLEU comparisons for distilled students on summarization and machine translation datasets.Source: Table 7 · p. 18

E.2. Full Results of Task-Specific KD (Tab. 3)

• DISTILLM consistently outperforms other KD baselines on XSum and CNN/DM in ROUGE-L and BLEU scores.The experiments subsample 20K training examples because existing student-generated-output methods can require up to 5× the training time of naïve KD.Source: §E.2 · p. 18

E.3. Terminal Probability by Adaptive SGO Scheduler

• Adaptive SGO terminal probabilities vary across model families, sizes, and datasets, supporting task-dependent balancing between student-generated outputs and fixed data.Reported terminal probabilities range from 0.2 for OPT-125M to 0.8 for T5-Base on SAMSum.Source: §E.3 · p. 18

E.4. Sensitivity Study for α

• ROUGE-L scores exhibit a concave pattern as α increases, linking the empirical sensitivity to the theoretical relationship between normalized approximation error and gradient scale.The reported interpretation identifies α as balancing approximation error against gradient scale.Source: §E.4 · p. 18, Table 8 · p. 18

E.5. Design of Replay Ratio

• The decreasing replay-ratio schedule achieves the highest training efficiency and the best overall performance, winning in three of five cases despite off-policy bias.Constant and increasing schedules perform similarly, which the authors attribute to adaptive SGO-use probabilities.Source: §E.5 · pp. 18–19, Table 9 · p. 19

E.6. Synergy of SKL and SRKL

• With α=0.1 and no SGO, ISKL outperforms previously used JSD and RKLD losses across all tested β values, although β has no consistent performance trend.ISKL interpolates between SKL and SRKL using β ∈ [0, 1].Source: §E.6 · p. 19, Figure 13 · p. 19

E.7. Qualitative Evaluation

• DISTILLM generates more detailed and accurate instruction-following responses, including correct alphabetical ordering, categorization, code generation, and mathematical reasoning.The qualitative comparisons use OpenLLaMA2-3B distilled models and prompts from Dolly, Self-Instruct, and Vicuna.Source: §E.7 · p. 19, Table 15 · p. 23

E.8. Replay Buffer Capacity

• A replay-buffer capacity of 1000 provides the most balanced performance, whereas smaller buffers risk overfitting and larger buffers retain outdated student-generated outputs that may increase bias.The capacity comparison is reported for GPT-2 students in the off-policy approach.Source: §E.8 · p. 19, Table 10 · p. 19