Multi-Token Prediction via Self-Distillation

John Kirchenbauer, Abhimanyu Hans, Brian R. Bartoldson, Micah Goldblum, Ashwinee Panda, Tom Goldstein

Keywords: efficient inference, multi token prediction, distillation

TL;DR: We accelerate LLM inference by adapting next-token prediction models for multi-token prediction using a self-distillation objective.

Abstract: Existing techniques for accelerating language model inference, such as speculative decoding, require training auxiliary speculator models and building and deploying complex inference pipelines. We consider a new approach for converting a pretrained autoregressive language model from a slow single next token prediction model into a fast standalone multi-token prediction model using a simple online distillation objective. The final model retains the exact same implementation as the pretrained initial checkpoint and is deployable without the addition of any auxiliary verifier or other specialized inference code. Our method produces models that decode more than $3\times$ faster at $<5\%$ drop in accuracy on GSM8K relative to the single token decoding performance of the same checkpoint.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 48