Issue #106 Training Edwin Dong, Anurag Shinde
On-Policy Distillation for LLMs, Simply Explained
Free editionPRO edition
This week’s PRO issue includes a Colab notebook, thoroughly dissecting on-policy distillation with intuitive Python code.
We’ll implement a simple teacher, try using supervised KD and on-policy KD, then generalize it to GKD while testing it using the interpolation parameters. We’ll finally pair GKD with RL and compare all KD methods. Please do read the full article before attempting the code tutorial.
As always, the code in this issue can be executed on a normal home PC and does not require a large amount of VRAM.
The monthly interview question bank can be found here. This month, we’ve incorporated more company-specific interview questions.
These are industry standard questions that are tested across major tech companies, AI research firms, and more. The questions are sorted by difficulty and topic.
Check back frequently for more questions, and let us know what type of questions you would like to see next!
Large language models are impressive. They summarize articles, translate languages, solve math problems, and write code.
However, one of the key problems surrounding them is that they are expensive. Serving a model with billions of parameters costs real money at real scale.
This has made model compression a critical topic in modern AI research.
One of the most effective compression strategies is knowledge distillation: train a smaller student model to replicate the behavior of a larger teacher model. The student runs faster, uses less memory, and can be deployed cheaply.
This is simple in principle but harder in practice.
The Problem With LLM Distillation
Here is the core issue with how distillation is done today.
During training, the student model learns from a fixed dataset of sequences. These sequences are either from the ground-truth data (teacher forcing) or generated by the teacher in advance.
During inference, the student generates its own sequences from scratch, one token at a time.
The problem, however, is that these are two completely different situations.
The sequences the student encounters while training look nothing like the sequences it will generate when deployed. More formally, the distribution of contexts seen during training differs from the distribution encountered at inference.
This gap is called the train-inference distribution mismatch, and it is a well-known problem in imitation learning research.
This edition is for PRO members.
