让大模型一次预测多个未来词,速度提升近5倍。
Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential
- 用掩码输入联合预测多个后续词,突破逐字生成限制。
- 代码和数学生成快近5倍,通用对话快2.5倍,质量不变。
- 适合追求推理加速的开发者和部署场景。
自回归语言模型因逐字生成的特性受限于推理速度与并行性,尤其在文本后期方向明确时更为明显。本文提出一种新框架,利用原始模型对未来的隐含知识,实现多个后续词的同时预测。方法包括:(1)通过掩码输入从共同前缀联合预测多个未来词;(2)采用门控LoRA保持原模型功能,同时支持多词预测;(3)引入轻量可学习采样模块,生成连贯序列;(4)设计一致性等辅助损失,提升生成一致性与准确率;(5)采用推测生成策略,使未来token扩展呈平方级增长,同时保持高保真度。经微调后,该方法在预训练模型上实现显著加速:代码与数学生成提速近5倍,通用对话与知识任务提速约2.5倍,且无质量损失。
原文摘要 · Abstract (English)
Autoregressive language models are constrained by their inherently sequential nature, generating one token at a time. This paradigm limits inference speed and parallelism, especially during later stages of generation when the direction and semantics of text are relatively certain. In this work, we propose a novel framework that leverages the inherent knowledge of vanilla autoregressive language models about future tokens, combining techniques to realize this potential and enable simultaneous prediction of multiple subsequent tokens. Our approach introduces several key innovations: (1) a masked-input formulation where multiple future tokens are jointly predicted from a common prefix; (2) a gated LoRA formulation that preserves the original LLM's functionality, while equipping it for multi-token prediction; (3) a lightweight, learnable sampler module that generates coherent sequences from the predicted future tokens; (4) a set of auxiliary training losses, including a consistency loss, to enhance the coherence and accuracy of jointly generated tokens; and (5) a speculative generation strategy that expands tokens quadratically in the future while maintaining high fidelity. Our method achieves significant speedups through supervised fine-tuning on pretrained models. For example, it generates code and math nearly 5x faster, and improves general chat and knowledge tasks by almost 2.5x. These gains come without any loss in quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。