Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
Authors: name: Haiduo Huang; name: Fuwei Yang; name: Zhenhua Liu; name: Yixing Xu; name: Jinze Li; name: Yang Liu; name: Xuanwu Yin; name: Dong Li; name: Pengju Ren; name: Emad Barsoum
Published: 2025-02-10 09:24:06

Abstract: 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.