Source-linked AI summary
Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE
Haiduo Huang, Fuwei Yang, Zhenhua Liu, Yixing Xu, Jinze Li, Yang Liu, Xuanwu Yin, Dong Li, Pengju Ren, Emad Barsoum
TL;DR
Speculative decoding can accelerate LLM inference, but candidates derived from the same representation limit prediction diversity. Jakiro uses MoE experts to decouple candidates and combines autoregressive and parallel decoding with feature-level contrastive operations; it reports improved accuracy and speedups across benchmarks and models.
Problem
Speculative-decoding candidates derived from the same representation remain correlated, limiting their diversity and draft prediction accuracy.
Source: §1 · p. 2Method
Jakiro uses MoE experts to decouple draft-tree candidates and combines autoregressive initial-token prediction with parallel decoding and feature-level contrastive operations.
Source: §1 · p. 2Results
Takeaways & Limitations
Across tasks and model sizes, Jakiro reports superior speedup ratios, high acceptance lengths, and efficiency without sacrificing model output quality.
Source: §6 · p. 9Jakiro has no substantial advantage on strictly top-1 accuracy tasks such as HumanEval and GSM8K, which are optimized for greedy decoding.
Source: §3.2 · p. 4
Abstract
from arXiv · showhide
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.
1. Introduction
Jakiro targets limited diversity among speculative-decoding candidates that share feature representations. It combines MoE-based decoupling with hybrid decoding and feature-level contrastive operations, and reports SOTA performance across benchmarks and target models.
1. Introduction
- Speculative decoding reduces inference latency by having a smaller draft model propose tokens for parallel verification by a larger target model.The approach is described as enabling lossless speedups by leveraging hardware parallelism during verification.Source: §1 · p. 1
- MoE experts generate draft-tree candidates to decouple their prediction correlations, increasing representation diversity with minimal computational cost.The approach is described as supporting both static and dynamic tree structures.Source: §1 · p. 2
- Tree-based candidates can remain correlated because they derive from the same representation, limiting draft prediction accuracy and token diversity.The introduction describes shared features among candidates within a layer as a remaining limitation even in Eagle-style approaches.Source: §1 · pp. 1–2
- Jakiro pairs autoregressive prediction for initial tokens with parallel decoding in the last two steps, adding contrastive operations to MoE dual-branch heads.The contrastive operation is between the two activated expert heads and is reported to add almost no latency.Source: §1 · p. 2
- Jakiro reports SOTA performance in speculative decoding across experiments on common models and benchmarks, with notable gains over prior SOTA on non-greedy MT-bench.Figure 2 covers Vicuna, LLaMA2-chat, and LLaMA3-instruct on MT-bench at temperature 1; its caption does not provide the speedup values.Source: §1 · p. 2, Figure 2 · p. 2
2. Preliminaries
Speculative decoding drafts candidate tokens, verifies them in parallel with a target model, and refines the output while preserving the target model's output distribution. Medusa-style heads share a target representation, while Eagle-style decoding decouples draft tokens across time steps but retains coupling within a tree layer.
2.1. Speculative Decoding
- Speculative decoding drafts candidate tokens from a prefix, verifies their probabilities in one target-model pass, then appends accepted tokens and resamples rejected ones.The passage states that the process preserves the output distribution of vanilla autoregressive decoding for greedy and non-greedy sampling.Source: §2.1 · p. 3
2.2. Medusa-style Decoding
- Medusa-style decoding uses multiple lightweight heads on the target model's last hidden state, but their shared representation limits diverse, independent predictions.The heads are typically small MLPs.Source: §2.2 · p. 3
2.3. Eagle-style Decoding
- Eagle-style decoding autoregresses over features to decouple draft tokens across time steps, but top-k tokens within a tree layer remain coupled.Eagle1 uses a static tree, while Eagle2 adds dynamically adjustable trees; neither removes the stated intra-layer coupling.Source: §2.3 · p. 3
3. Jakiro
Jakiro uses dynamically allocated MoE experts to decouple draft-token predictions and diversify tree candidates while preserving inter-layer independence. It combines this design with contrastive inference and parallel decoding, and trains the draft model with a combined objective.
3.1.1. DYNAMIC DECOUPLING WITH MOE HEADS
- Jakiro uses a single decoder layer with parallel lightweight expert heads, dynamically allocating experts to decouple token predictions while retaining the target model’s embedding and head layers.Unlike Medusa, which uses the last hidden state for all draft heads, Jakiro dynamically allocates expert modules across tokens.Source: §3.1.1 · p. 3
3.1.2. MOE TREE CONSTRUCTION
- Jakiro’s MoE tree decouples intra-layer dependencies while preserving inter-layer independence, broadening candidate diversity and improving prediction accuracy and factuality, especially in non-greedy modes.The approach can be combined with either Eagle1’s static tree or Eagle2’s dynamic tree.Source: §3.1.2 · pp. 3–4
3.2. Efficient Integration of Contrastive Mechanism
- Contrastive inference compares the two highest-scoring experts’ output hidden states, not logits, while parallel decoding is focused on the final two steps.Jakiro does not show a substantial advantage on strictly top-1 tasks; parallel decoding for K > 2 is theoretically supported but left for future work.Source: §3.2 · pp. 4–5
3.3. Training of the Draft Models
- Draft training uses a fixed preprocessed dataset, Smooth L1 feature regression and token cross-entropy in a combined loss.The parallel-decoding mechanism at the penultimate step adds a loss based on contrast between the two MoE expert heads.Source: §3.3 · p. 5
4. Experiments
Experiments evaluate Jakiro across mainstream model sizes and varied tasks, reporting speedup and average acceptance length. Ablations examine expert counts and the effects of parallel decoding and contrastive features.
4. Experiments
- Evaluation spans Vicuna, LLaMA2-Chat, and LLaMA3-Instruct sizes across dialogue, code, math, instruction following, summarization, and question answering.Tests cover model sizes from 7B to 70B and use several GPU types and multi-GPU configurations for larger models.Source: §4 · pp. 5–6
- Most experiments use batch size 1; metrics are end-to-end walltime speedup and average tokens accepted per target-model forward pass.The authors consider output-quality evaluation unnecessary because Jakiro maintains the target model’s output distribution.Source: §4 · p. 6
- Jakiro keeps target LLMs fixed and trains on ShareGPT for 68,000 dialogue iterations, with draft-model parameter counts from 0.35B to 1.87B.Training the 70B draft model takes approximately 2–3 days on 4×A100 40GB GPUs.Source: §4 · p. 6
4.1. Effectiveness
- 15.4% higher final average speedup than Eagle2 (2.84 vs. 2.46) across benchmark test sets for Vicuna 7B in non-greedy mode.For Vicuna 7B in greedy mode, speedups were 2.99x on MT-bench and 3.43x on HumanEval; GSM8K reached 3.11x with a 4.95-token average acceptance length.Source: §4.1 · p. 6, Table 1 · p. 7
4.2. Ablation Study
- Increasing candidate experts raises both average acceptance length and computational load, making the N-K setting a speed trade-off.The ablation varies candidate experts N and activated experts K on Vicuna 7B under MI250; tree construction can activate K≥2 experts per inference.Source: §4.2 · pp. 6–7, Table 3 · p. 7
- Parallel decoding and contrastive features each improve over MoE alone; combining them further improves speedup while keeping acceptance rate comparable.The ablation tests whether penultimate draft inference uses parallel decoding and feature contrast.Source: §4.2 · p. 8, Table 4 · p. 8
5. Related Work
Prior speculative-decoding work introduced tree attention, multiple draft heads, and autoregressive drafting, but limitations in token diversity and alignment remain. Jakiro combines MoE-based decoupling with parallel decoding and contrastive mechanisms.
Speculative decoding
- Medusa’s shared final hidden state limits independent draft-head predictions, while Eagle-style methods still couple top-k tokens within a draft-tree layer.Jakiro targets this same-layer coupling with dynamic decoupling using MoE heads.Source: §5 · p. 8
Mixture of Experts
- MoE heads are used for specialization in language and multimodal models, and Jakiro applies them to draft tokens to increase diversity and prediction confidence.The passage situates MoE among prior work including GShard (Lepikhin et al., 2021), Switch Transformer (Fedus et al., 2021), and StableMoE (Dai et al., 2022).Source: §5 · p. 8
Parallel decoding
- Prior parallel-decoding methods face draft-target alignment challenges; Jakiro aims to improve alignment and token diversity through decoupled MoE heads.The related-work discussion cites Monea et al. (2023) and Yi et al. (2024) as applying parallel decoding in speculative frameworks.Source: §5 · p. 8
6. Conclusion
Jakiro addresses speculative decoding’s challenges in token diversity and draft-target distribution alignment with dynamic MoE decoupling and contrastive mechanisms. The authors report strong speedup and acceptance length across tasks and model sizes in greedy and non-greedy generation.
6. Conclusion
- Jakiro’s dynamic MoE-head decoupling and contrastive mechanism target token diversity, prediction confidence, and draft-target distribution alignment.The conclusion frames these as remaining challenges for speculative decoding.Source: §6 · p. 8
- Across varied tasks and model sizes, Jakiro reports superior speedup ratios, high acceptance lengths, and efficiency without sacrificing model output quality.The reported scope includes both greedy and non-greedy generation tasks.Source: §6 · p. 9
A.1. Average Acceptance Length.
Figure 6 compares average acceptance length across Vicuna, LLaMA2-chat, and LLaMA3-instruct models on MT-bench in non-greedy settings. Jakiro has the highest acceptance length across the models, according to the text.
A.1. Average Acceptance Length.
- Jakiro achieves the highest average acceptance length across the compared models in non-greedy MT-bench inference.The plotted results are averaged over four inference runs on an A100-40G GPU.Source: §A.1 · p. 12, Figure 6 · p. 12
A.2. Results of Other Devices.
Jakiro shows consistent speedup improvements across devices, with Figure 7 reporting Vicuna 7B results across tasks. The figure distinguishes greedy and non-greedy sampling and averages four inference runs.
A.2. Results of Other Devices.
- The Figure 7 reproduction compares speculative-sampling methods that do not finetune backbone models, preserving the output text distribution.Results are averaged over four inference runs; T0 denotes greedy sampling and T1 non-greedy sampling.Source: Figure 7 · p. 12
- Jakiro achieves consistent speedup improvements across different devices, indicating performance transferability.Figure 7 covers Vicuna 7B inference across tasks.Source: §A.2 · p. 12, Figure 7 · p. 12
- Tables 5 and 6 report speedup ratios and average acceptance lengths, distinguishing MoE-based weighted decoupling from weighted non-decoupling with contrastive decoding.Medusa-style methods are excluded because their relaxed non-greedy acceptance conditions do not guarantee lossless acceleration.Source: Table 5 · p. 13, Table 6 · p. 14