⏶15
OmniDraft：一种用于设备端推测解码的跨词汇在线自适应起草器
发表
由 
Zhaocong Yuan 提交
Zhaocong Yuan 提交作者: 
Ramchalam Kinattinkara Ramakrishnan, Zhaocong Yuan, Shaojie Zhuo, 
Chen Feng, Yicheng Lin, 
Chenzheng Su, Xiaopeng Zhang
Ramchalam Kinattinkara Ramakrishnan, Zhaocong Yuan, Shaojie Zhuo, Chen Feng, Yicheng Lin, Chenzheng Su, Xiaopeng Zhang摘要
推测解码通常要求有一个小型、高效的草稿模型，该模型要么经过预训练，要么离线蒸馏到特定的目标模型系列，例如Llama或Qwen模型。然而，在在线部署设置中存在两个主要挑战：1）目标模型与草稿模型不兼容；2）期望随着使用和时间的推移，延迟有所改善。在这项工作中，我们提出了OmniDraft，一个统一的框架，它使单个草稿模型能够与任何目标模型配合使用，并动态适应用户数据。我们引入了一个在线n-gram缓存和混合蒸馏微调，以解决草稿模型和目标模型之间的跨词汇不匹配问题；并通过利用自适应草稿技术进一步提高解码速度。OmniDraft特别适用于设备端LLM应用，在这些应用中，模型成本、效率和用户定制是主要的争议点。这进一步强调了解决上述挑战的必要性，并推动了“一草稿模型通吃所有”的范式。我们通过在数学推理、编码和文本生成任务上执行在线学习，展示了OmniDraft框架的熟练度。值得注意的是，OmniDraft使单个Llama-68M模型能够与包括Vicuna-7B、Qwen2-7B和Llama3-8B模型在内的各种目标模型配对进行推测解码；此外，还提供了高达1.5-2倍的加速。
我们提出了 OmniDraft，它使一个草稿模型能够与任何目标模型（Llama、Qwen 等）在推测解码中协同工作。为了克服词汇不匹配问题，我们引入了跨词汇推测解码和蒸馏，并使用在线构建的 n-gram 缓存。我们还利用自适应草稿生成以进一步提高接受率和加速。总的来说，我们探索了“一拖多”的范式，为推测解码的更灵活设计开辟了道路。