Knowledge distillation from FP32 teacher to INT4 student using on-policy (student-generated) data rather than static datasets.

Standard KD trains on fixed datasets, causing distribution mismatch—the student never learns from its own mistakes. On-policy distillation fixes this by having the student generate sequences and learning from teacher feedback on those generations. Inspired by GKD and on-policy distillation.

On-policy distillation doesn't provide significant benefits over off-policy, even after hyperparameter tuning. Best-tuned λ=0 and λ=1 perform within 0.3% of each other.

• Distribution mismatch is probably not a big issue when student and teacher share the same weights at different precision
• On-policy rollouts are much more compute intensive than simple off-policy SFT
• One interesting finding: optimal β differs by regime (forward KL for off-policy, reverse KL for on-policy)

uv sync

# Off-policy KD (λ=0, β=0, lr=5e-6) - best tuned
uv run accelerate launch train.py @configs/offpolicy.toml --output-dir dump/offpolicy --quant-type int4

# On-policy KD (λ=1, β=1, lr=5e-6) - best tuned
uv run accelerate launch train.py @configs/onpolicy.toml --output-dir dump/onpolicy --quant-type int4

λ controls data source interpolation:

• λ=0: Off-policy (dataset sequences only)
• λ=1: Fully on-policy (student-generated sequences only)

Distillation recovers most of the accuracy lost by naive PTQ INT4, but on-policy (λ=1) and off-policy (λ=0) perform equivalently. The ~0.5% differences between them are within noise.

All within ~0.3% of each other—confirms null result. Note: optimal β differs by regime (forward KL for off-policy, reverse KL for on-policy).

Higher learning rates achieve lower training loss but worse eval accuracy. Lower LRs (5e-6, 1e-5) generalize better despite not fitting the training data as tightly.

Same pattern as off-policy: lower LRs (5e-6, 1e-5) generalize better despite higher training loss.

20x compute reduces training loss by 42% but eval accuracy is mixed. Suggests we've hit the accuracy ceiling for INT4 quantization.

• beta=0: forward KL
• beta=1: reverse KL
• beta between 0-1 interpolates between the two.

For standard distillation, forward KL is often used, however on policy distillation recommends reverse KL. I find similar results for λ=0 (off-policy distillation) forward KL performs better.

Results are close, but β=1 (reverse KL) is theoretically motivated for on-policy—the student should mode-seek rather than mean-seek when learning from its own generations. This aligns with on-policy distillation recommendations.

Interestingly, optimal β differs by training regime:

• Off-policy (λ=0): Forward KL (β=0) works best
• On-policy (λ=1): Reverse KL (β=1) works best

This makes theoretical sense: off-policy trains on fixed data (mean-seeking appropriate), while on-policy trains on student generations (mode-seeking appropriate).

From limited sweeps, it seems that batch size and rollout length are not sensitive params.

Note I do the same number of steps, so doubling batch size / rollout lenght doubles the compute budget! Suggesting we've saturated the quantization accuracy of the model.

Interpolating between on-policy and off-policy data sources.

Mixed strategy performs similarly to both extremes—no benefit from interpolation.