Reproduce “DistillSpec: Improving Speculative Decoding via Knowledge Distillation” on OpenResearch

About this paper

Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, Rishabh Agarwal

DistillSpec, developed by researchers at Google Research and Google DeepMind, enhances speculative decoding for large language models by applying knowledge distillation to train a smaller draft model for better alignment with a larger target model. The method significantly improves inference speed, delivering 10-45% faster generation across various benchmarks, and shows strong transferability to unseen tasks.

Key insights

• Training distillation models with on-policy data generated by the draft model itself is effective and computationally efficient for improving speculative decoding acceptance rates.
• White-box knowledge distillation, leveraging teacher model logits, is essential for achieving the fine-grained alignment needed between draft and target models for optimal speculative decoding performance.
• The most effective divergence function for knowledge distillation in the context of speculative decoding varies depending on the specific task and decoding strategy (e.g., greedy vs. non-greedy sampling).

Summary by alphaXiv.