用专家模型提升大模型推理速度,预测更准且更快。
Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
- 引入专家模型分离生成候选,打破预测相关性。
- 混合解码+对比特征增强,准确率显著提升。
- 适合追求高速高精度推理的模型部署场景。
推测解码(SD)通过小型草稿模型预测多个词元,并由大型目标模型并行验证,从而加速大语言模型推理。然而,草稿模型容量有限,常需采用树形采样以提高预测准确率,即每步生成多个候选。我们发现该方法的关键局限在于:同一时刻的候选均来自相同表示,限制了多样性,降低整体效果。为此,我们提出Jakiro,利用专家混合(MoE)机制,使独立专家生成多样化预测,有效解耦候选间的相关性。此外,我们设计混合推理策略,初始阶段采用自回归解码,后续阶段采用并行解码,并在后者中引入特征对比机制以提升准确性。实验表明,该方法显著提高预测准确率,实现更高推理加速比。在多种模型上的广泛测试验证了其有效性与鲁棒性,确立了推测解码新SOTA。代码已公开于https://github.com/haiduo/Jakiro。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。