GKD · On-policy Distillation

On-policy Distillation of Language Models: Learning from Self-Generated Mistakes

Rishabh Agarwal et al. · ICLR 2024

Read paper · arXiv:2306.13649v3

Query the teacher on student-generated prefixes to reduce the mismatch between training and generation.

What you will implement

Implement student rollouts and full-vocabulary divergence in a tiny autoregressive model. OPD here means on-policy distillation.

• 12. On-policy Distillation at Student Prefixes

  Sample from the student and distill full-vocabulary feedback at the same prefixes.

  GKD with lambda=1 and explicit forward/reverse KL; a supplied three-token Markov LM replaces large language models.