arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2502.06282cs.CLcs.AIcs.LG

Jakiro:通过MoE解耦多头提升投机解码

Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE

Advanced Micro Devices, Inc.(超威半导体公司);Institute of Artificial Intelligence and Robotics, Xi'an Jiaotong University(西安交通大学人工智能与机器人研究所);University of HongKong(香港大学)

Haiduo Huang, Fuwei Yang, Zhenhua Liu, Yixing Xu, Jinze Li, Yang Liu, Xuanwu Yin, Dong Li, Pengju Ren, Emad Barsoum

AI总结:

Jakiro利用MoE解耦候选生成提升多样性,结合自回归与并行混合推理及对比机制,显著提高投机解码准确性与加速比,达到新SOTA。

摘要:

投机解码(SD)通过使用较小的草稿模型预测多个token,然后由较大的目标模型并行验证,从而加速大语言模型推理。然而,草稿模型的有限容量通常需要基于树的采样来提高预测准确性,即在每一步生成多个候选。我们发现该方法的一个关键局限:同一步的候选来自相同的表示,限制了多样性并降低了整体效果。为解决这一问题,我们提出Jakiro,利用专家混合(MoE),其中独立的专家生成多样化的预测,有效解耦候选之间的相关性。此外,我们引入一种混合推理策略,将初始token的自回归解码与后续阶段的并行解码相结合,并通过特征中的对比机制增强后者以提高准确性。我们的方法显著提升了预测准确性,并实现了更高的推理加速。跨多种模型的广泛实验验证了我们方法的有效性和鲁棒性,在投机解码中建立了新的SOTA。我们的代码可在https://github.com/haiduo/Jakiro获取。

英文摘要:

Speculative decoding (SD) accelerates large language model inference by using a smaller draft model to predict multiple tokens, which are then verified in parallel by the larger target model. However, the limited capacity of the draft model often necessitates tree-based sampling to improve prediction accuracy, where multiple candidates are generated at each step. We identify a key limitation in this approach: the candidates at the same step are derived from the same representation, limiting diversity and reducing overall effectiveness. To address this, we propose Jakiro, leveraging Mixture of Experts (MoE), where independent experts generate diverse predictions, effectively decoupling correlations among candidates. Furthermore, we introduce a hybrid inference strategy, combining autoregressive decoding for initial tokens with parallel decoding for subsequent stages, and enhance the latter with contrastive mechanism in features to improve accuracy. Our method significantly boosts prediction accuracy and achieves higher inference speedups. Extensive experiments across diverse models validate the effectiveness and robustness of our approach, establishing a new SOTA in speculative decoding. Our codes are available at https://github.com/haiduo/Jakiro.

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

↑