arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.06019cs.CLcs.LG
收藏

通过自蒸馏实现多token预测

Multi-Token Prediction via Self-Distillation

University of Maryland(马里兰大学);Lawrence Livermore National Laboratory(劳伦斯利弗莫尔国家实验室);Columbia University(哥伦比亚大学);TogetherAI

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

AI总结:

本文提出通过自蒸馏将预训练的自回归语言模型转换为快速独立的多token预测模型,无需额外辅助模型或复杂管道,实现更快的解码速度和更高的准确性。

摘要:

现有加速语言模型推理的技术,如推测解码,需要训练辅助推测模型并构建复杂的推理管道。我们考虑了一种新的方法,通过简单的在线蒸馏目标,将预训练的自回归语言模型从慢速的单token预测模型转换为快速的独立多token预测模型。最终模型与预训练的初始检查点保持完全相同的实现,并且无需添加任何辅助验证器或其他专用推理代码即可部署。我们的方法在GSM8K数据集上,相对于相同检查点的单token解码性能,解码速度超过3倍,准确率下降不到5%。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑