Jakiro：通过 MoE 实现解耦式多头提升推测解码
计算与语言
2025-03-12 v1 人工智能
机器学习
摘要
推测解码 (SD) 通过使用较小的草稿模型预测多个标记，然后由较大的目标模型并行验证来加速大语言模型的推理。然而，草稿模型的有限容量常常需要基于树的抽样来提高预测准确度，即在每一步生成多个候选选项。我们识别出这一方法的一个关键限制：同一步骤上的候选选项源自相同的表示，限制了多样性，降低了整体效果。为此，我们提出 Jakiro，利用 Mixture of Experts (MoE)，其中独立的 expert 生成多样化的预测，有效地解耦了候选选项之间的相关性。此外，我们引入一种混合推理策略，将初始标记的自回归解码与后续阶段的并行解码相结合，并通过对特征进行对比机制来增强后者的准确度。我们的方法显著提高了预测准确度，并实现了更高的推理加速。广泛的实验表明，我们的方法在多样模型上有效且稳健，建立了推测解码的新的 SOTA。我们的代码已公开：https://github.com/haiduo/Jakiro。
引用
@article{arxiv.2502.06282,
title = {Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE},
author = {Haiduo Huang and Fuwei Yang and Zhenhua Liu and Yixing Xu and Jinze Li and Yang Liu and Xuanwu Yin and Dong Li and Pengju Ren and Emad Barsoum},
journal= {arXiv preprint arXiv:2502.06282},
year = {2025}
}