projects / 03
distill-lab
knowledge distillation, four ways
Qwen3-0.6B-Base learns grade-school math (GSM8K) from Qwen3-4B four ways: fine-tuning on human-written solutions as a control, fine-tuning on the teacher's own solutions, and matching the teacher's token distributions with forward and with reverse KL. Every student gets the same evals on GSM8K and on SVAMP, a set of problems written in a different style. The losses and the training loop are written by hand in PyTorch, and every run fits on a free Kaggle T4.
- GSM8K over base (49.4%)
- +15 pts
- GSM8K over base (49.4%)
- SVAMP, reverse KL (text: 62.6%)
- 70.3%
- SVAMP, reverse KL (text: 62.6%)
- logit cache (full: ~320 GB)
- 202 MB
- logit cache (full: ~320 GB)
highlights
- Matching the teacher's token probabilities beat copying its text on unfamiliar problems: 70.3% vs 62.6% on SVAMP
- Training on the teacher's solutions added about 15 points on GSM8K; human-written solutions added nothing
- Top-32 teacher logits cached in 202 MB instead of ~320 GB, keeping 99.99% of the teacher's probability
- Losses and training loop hand-written in PyTorch; each run trains in under 30 minutes on one T4
how it works
- 01
Baseline
Student and teacher answer GSM8K's 1,319 test questions and SVAMP's 1,000 with a 4-shot prompt and greedy decoding: 49.4% and 63.0% for the student, 85.1% and 84.0% for the teacher. Every answer is extracted programmatically, and the fallback cases were checked by hand.
- 02
Teacher data
The teacher writes two solutions for each of the 7,465 training questions and only correct ones are kept: 12,724 after removing duplicates, none overlapping a test question.
- 03
Cache the teacher's logits
The teacher reads its kept solutions back, and its top-32 next-token log-probs are saved at all 1,050,428 solution positions, so logit distillation never has to run the 4B model during training.
- 04
Train
Full fine-tune for one epoch, learning rate 1e-5 with warmup and cosine decay, 32 sequences per batch. S1 trains on human solutions and R1 on the teacher's text; L1 and L2 match the teacher's distribution with 0.9 × KL (forward and reverse) plus 0.1 × cross-entropy on the text.
- 05
Evaluate
Every student runs the same zero-shot evals on GSM8K and SVAMP, reported with 95% binomial intervals so differences inside the noise aren't read as results.
design decisions
Top-32 logits over the full distribution
Storing all 151,936 log-probs per position would take about 320 GB; the top 32 take 202 MB and keep 99.99% of the teacher's probability. Measured before training: forward KL loses 0.04%, reverse KL comes out about 6% low because the student's mass on the dropped tokens lands in one bucket.
Loss only on graded tokens
The first L1 run ran out of memory on the T4: every loss did a full 152k-vocabulary fp32 softmax at every padded position, three times per batch. Computing it only on the ~82 graded tokens per example fixed it, and L1 now trains as fast as R1.
A human-solution control
S1 trains on the dataset's own solutions with the same settings. It stays within noise of the base model, so the gains in R1, L1 and L2 come from the teacher, not from more fine-tuning on GSM8K.
known limitations
- Every result is a single run, so the 2.3-point gap between forward and reverse KL is inside the noise.
- R1 saw more sequences than S1 (up to two solutions per question); the data-size runs that would separate volume from source weren't done.
- The students never see the 7.5% of questions the teacher couldn't solve.
- Calibration and on-policy distillation were planned but not run.
stack
- python
- pytorch
- hugging face
- qwen3
- kaggle