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

Knowledge distillation trains a small student to match a large teacher's next-token distributions. Two choices matter: which loss (forward KL, reverse KL, or blends) and which data the loss is computed on (responses written by the teacher, or by the student itself). Forward KL rewards covering everything the teacher might say, which in a small student smears probability over bad outputs; reverse KL rewards saying only what the teacher likes, which can make the student narrow. Student-written data closes the gap between training and inference, but is lower quality.

DistiLLM-2 pairs each loss with the data it suits. On teacher responses, a skew forward KL (KL from the teacher to a mix of teacher and student) pulls the student's probability up on what the teacher writes. On student responses, a skew reverse KL pushes the student's probability down where the teacher assigns little. The two terms are weighted (1 − β) and β, with β rising linearly from 0.5 toward 1 during training. The skew coefficient α is set per sample from how far the student's probability is from the teacher's. Student responses are regenerated once per epoch rather than every step, which works with fast batched inference.

The authors also show what does not work. Plugging DPO's preference loss in with the teacher as the reference model (DPKD) maximizes the gap between the two responses and drives the student into degenerate text. Using responses from stronger models than the teacher on the forward-KL side did worse than the teacher's own text.

DistiLLM-2 beat KD, SeqKD, ImitKD, GKD, DistiLLM and Speculative KD on instruction following for three teacher-student pairs, on GSM8K and MATH, and on HumanEval and MBPP, and helped as a reference model for DPO, for vision-language students, for recovering INT4 quantized models and for training speculative-decoding drafters.

Customers who settle on a capable 7B to 9B model often want something smaller for edge devices, latency or cost. Distillation is the standard path, and the choices this paper studies are the ones a team makes, usually by default: which loss, on which data.

The default is often the weakest. Plain KD (forward KL on teacher outputs) scored 41.19 average win rate for Qwen2-1.5B, against 56 to 59 for methods that also use student outputs. If a team's distillation script only replays teacher text, it is leaving most of the gain behind.

Batched student generation is cheap enough. Regenerating student responses once per epoch with a fast inference engine matched fully on-policy training on the authors' Dolly test (ROUGE-L 26.20 vs 26.37), so the on-policy benefit does not require a slow training loop.

Drafters for speculative decoding. A tiny draft model trained this way was accepted more often by the large model, giving 1.97x speedup on Phi-3-medium. If a customer serves a large model and wants lower latency, distilling a better drafter is a direct application.

The gains are a few points, not a new tier. Win-rate gains over the next best method were 2.0 to 4.5 points; math gains over DistiLLM were 0.2 to 1.8. The student still trails the teacher by 15 to 27 points of win rate.