Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866909533464428544 |
|---|---|
| author | Huang, Haiduo Yang, Fuwei Liu, Zhenhua Xu, Yixing Li, Jinze Liu, Yang Yin, Xuanwu Li, Dong Ren, Pengju Barsoum, Emad |
| author_facet | Huang, Haiduo Yang, Fuwei Liu, Zhenhua Xu, Yixing Li, Jinze Liu, Yang Yin, Xuanwu Li, Dong Ren, Pengju Barsoum, Emad |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_06282 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE Huang, Haiduo Yang, Fuwei Liu, Zhenhua Xu, Yixing Li, Jinze Liu, Yang Yin, Xuanwu Li, Dong Ren, Pengju Barsoum, Emad Computation and Language Artificial Intelligence Machine Learning 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. |
| title | Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE |
| topic | Computation and Language Artificial Intelligence Machine Learning |
| url | https://arxiv.org/abs/2502.06282 |