Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
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.
67ab6c9867a1607ab478b975
Haiduo Huang
claimed_verified
2025-02-11T15:33:50.773Z
65780c60411e14898b8da93e
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/fvjP6qICiuR09LV8Xzahb.png
Haiduo Huang
user
Hhaiduo
67ab6c9867a1607ab478b976
Fuwei Yang
67ab6c9867a1607ab478b977
Zhenhua Liu
67ab6c9867a1607ab478b978
Yixing Xu
67ab6c9867a1607ab478b979
Jinze Li
67ab6c9867a1607ab478b97a
Yang Liu
67ab6c9867a1607ab478b97b
Xuanwu Yin
67ab6c9867a1607ab478b97c
Dong Li
67ab6c9867a1607ab478b97d
Pengju Ren
67ab6c9867a1607ab478b97e
Emad Barsoum
Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
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.
Jakiro, a speculative decoding method using Mixture of Experts, enhances prediction accuracy and inference speed by generating diverse candidates and combining autoregressive and parallel decoding strategies.
speculative decoding draft model target model tree-based sampling Mixture of Experts independent experts hybrid inference strategy autoregressive decoding parallel decoding contrastive mechanism