Jakiro: Boosting Speculative Decoding with Decoupled Multi-Head via MoE

C1: This claim attributes speedup to "minimizing the risk of errors during the inference process". What errors are the authors talking about?

Thank you for pointing out the ambiguity in our phrasing. The "risk of errors" refers to the probability that a token generated by the draft model is rejected by the target model during the validation phase of Speculative Decoding (SD). This rejection risk directly corresponds to the inverse of the acceptance rate of the draft model. When a draft token is rejected at a particular step:

• All subsequent tokens in the speculative sequence are discarded.
• The system must re-sample a new token from the target model, which resets the speculative decoding process for subsequent steps.

A low acceptance rate (i.e., high rejection risk) leads to frequent re-sampling events, which disrupts the efficiency of SD. This is because:

• Computational Overhead: Re-sampling from the target model incurs additional latency, compromising the speed gains from speculative decoding.
• Loss of Speculative Progress: Discarded speculative tokens waste the computational effort invested in generating them.

Jakiro’s design mitigates this risk by optimizing the draft model’s alignment with the target model, thereby increasing the acceptance rate. This results in fewer re-sampling events and more stable, longer sequences of accepted tokens, ultimately enhancing the overall speedup of the SD process.

Q1: The work does not talk about the inference framework they used or if they did or did not use chunked prefill. Overall, the paper presents a promising approach to speculative decoding with supporting evidence.

Thanks for your valuable comment. Similar to mainstream SD methods (e.g., Medusa, EAGLE, Hydra), our Jakiro implementation relies solely on the PyTorch framework without additional acceleration architectures.

Regarding "chunked prefill" (Agrawal et al., 2023), we didn't employ this technique because:

• Our experiments used batch_size=1 by default, may refer to our response to Reviewer aZSm's comment W2.
• The models could be fully loaded onto GPUs without memory constraints.

We acknowledge that integrating "chunked prefill" could benefit larger models or batch_size>1 scenarios. We plan to explore such optimizations (e.g., vLLM/SGLang integration) in future work to further enhance Jakiro's performance.

W1: The result seems very incremental. They used an MoE (which most people should be doing now anyway) and got slightly better results at the cost of needing to store more draft model parameters.

We appreciate this critical perspective but wish to clarify that characterizing our results as merely incremental may not be entirely accurate. While we do employ MoE, simply increasing model parameters does not guarantee improved speedup—expanding the draft model might improve acceptance rates but also introduces additional overhead, potentially reducing the speedup gain (as demonstrated in our ablation studies in Table 3).

Furthermore, we optimize the MoE architecture by using slimmer MLP dimensions (detailed in the lower part of Figure 3 in this paper). To verify that Jakiro introduces minimal additional parameters and almost no extra memory overhead during inference, please refer to our response to Reviewer LwQj’s comment W2.

We maintain that integrating MoE into speculative decoding remains an innovative and non-trivial contribution.

C1: The authors should give a clear definition of the diversity of draft tokens and then quantify it with such a definition.

Thanks for your thoughtful feedback. Building upon our previous response to Reviewer LwQj (W1), as illustrated in the figure, our Jakiro model, even with just two MOE heads, can generate richer tokens.

To rigorously address your query on defining and quantifying diversity, we propose a composite metric that evaluates three dimensions:

• Generation Richness (G): The number of unique draft tokens, used to evaluate the diversity of the drafting phase.

• Selection Effectiveness (S): The number of unique accepted tokens, used to evaluate the diversity of the verification phase.

• Exploration Depth (E): The total number of tokens generated in the final output, prioritizing deeper exploration.

The diversity metric is formalized as:

$$\text{Diversity} = G + S +E = \frac {\sum ^{i}_{N} {(\alpha \cdot \ln{g_i} + \beta \cdot \ln{s_i} + \gamma \cdot \ln{e_i} )}} {N}$$

where $g_i$, $s_i$, and $e_i$ represent the Generation Richness, Selection Effectiveness, and Exploration Depth for each round of dialogue, respectively, and $N$ denotes the total number of dialogue rounds. The empirically validated weights are: $\alpha=0.4,\ \beta=0.4,\ \gamma=0.2$.

Benchmark results on mt_bench (Vicuna-7B, T=1, top-k=10):

Jakiro achieves 11% higher diversity than the Eagle, driven by its ability to generate more unique tokens, maintain high acceptance rates, and explore longer sequences.

C2: The motivation for contrastively decoding with top-2 selected experts. Consider using the features from the most activated expert and least activated experts to do contrastive decoding.

Thank you for raising this important point. The rationale for using top-2 activated experts in contrastive decoding:

• Baseline Configuration: Our initial implementation followed the mainstream MoE framework, where the top-2 experts (ranked by routing weights) are selected for token generation.

• Empirical Optimization: As shown in Table 3 of the paper, using 2 experts achieved the optimal speedup.

To rigorously evaluate our design, we conduct additional Vicuna 7B experiments comparing two strategies on A40:

We fully agree with the reviewer on the theoretical significance of exploring "strong-weak expert contrastive decoding." However, our experiments reveal that low-confidence experts (e.g., bottom-1) may introduce noise, leading to degraded output quality. We analyze that this discrepancy from the original contrastive decoding conclusions (which compare logits from strong/weak models) might stem from: our contrastive decoding mechanism compares hidden states of experts. This architectural difference could explain why the "strong-weak expert" paradigm behaves differently in our framework.

R1: Missed some of the recent works in parallel SPD.

We will discuss recent parallel SPD works like BiTA, Parallel Decoding via Hidden Transfer, and ParallelSpec in the revised manuscript.

W1: The speed-up ratio suffers as the number of experts increases in Table 3. N=K=2 means a simple EAGLE-style implementation with two heads.

We would like to clarify two key points:

1. It fundamentally differs from EAGLE in that our Jakiro uses dynamic router-based expert selection for autoregressive phases (Stages 1-4) and employs contrastive decoding only in the final parallel phase (Stages 5-6).
2. In smaller models (e.g., Vicuna-7B), increasing N introduces additional computational costs (e.g., router computation), which slows down the autoregressive phase.

To address these trade-offs, we propose the following directions:

• Efficiency Optimization for Small Models: Explore lightweight routing mechanisms or parameter-sharing techniques to reduce N’s overhead.
• Dynamic N/K Adjustment: Adaptively set N and K based on task complexity (e.g., higher N for complex tasks, lower N for simplicity).

Our MoE design inherently supports scalable improvements for large models (e.g., DeepSeek-V3-671B), and the draft model’s overhead becomes negligible relative to total computation (Experts N>2).

W1: The authors should test over more advanced settings like flash decoding.

Thank you for the feedback. We first clarify that Flash Decoding (FD) and speculative decoding (SD) operate at different optimization levels but can be effectively combined for greater efficiency.

Flash Decoding (Dao et al., 2023):

A system-level optimization that accelerates attention computation for long-context processing through:

• Memory-efficient tiling for KV caching
• Parallel reduction across sequence chunks

Speculative Decoding (Leviathan et al., 2023):

An algorithmic innovation leveraging the computation-memory gap in LLM inference:

• Draft model proposes candidate tokens (reducing generation steps)
• Target model verifies in parallel (exploiting GPU underutilization)

By conducting Vicuna 7B experiments on A100 under T=0 (Refer to FlashDecoding++, Hong et al., 2023):

W2: The authors should test over throughput with different bsz.

Similarly, we conducted experiments on "Batch Sizes > 1" (T=0) with Vicuna 7B on A40, and the results are shown in the table below. These results validate that the speedup of our method increases as the batch size grows, but the gains diminish gradually.

Additional Clarification on Applying SD to Batch Sizes > 1:

• Performance Degradation with Larger Batches: Simply increasing the batch size during SD inference shifts the problem nature from memory-bound to compute-bound, leading to diminishing returns or even negative impacts as batch sizes grow (MagicDec, Sadhukhan et al., 2025).

• Practical Focus on Batch Size = 1: Current SD optimizations prioritize batch size = 1 due to

  • Sequence Length Variability: Divergent acceptance rates across sequences in a batch result in varying candidate sequence lengths after the drafting stage.

  • Verification Overhead: This inconsistency increases computational costs (or latency) during the parallel verification phase of the target model.

Research Status & Outlook: Existing SD implementations (e.g., the seminal SD work, Medusa, EAGLE, Hydra) predominantly target batch size = 1. Though optimizing for batch sizes > 1 remains an open challenge, it represents a promising direction for future research.

W3: If the improvement of 'contrastive loss' is marginal, the authors can omit this part to keep the method simple.

We appreciate the reviewer's observation regarding the contrastive loss. While the absolute improvement may appear modest, the contrastive loss is architecturally indispensable for maintaining the system's end-to-end performance. Specifically on GSM8K of Table 4 in this paper, the contrastive loss contributes to a measurable Speedup enhancement (3.05x → 3.11x). Its critical role lies in:

• Optimizing later-stage parallel decoding efficiency
• Preserving the acceptance rate during initial auto-regressive phases

Of course, this part can be omitted to maintain method simplicity if a certain performance loss is acceptable.

S1: The authors should also add models like Qwen with better at math and coding.

Thanks for your valuable suggestion. We have added comprehensive experiments on Qwen2-7B-Instruct in A40 (Speedup: λ, Average accepted length: τ):

These additions highlight Jakiro’s adaptability to specialized LLMs.

W1: The analysis of diversity is missing in the experiments.

To validate that our Jakiro method enhances the diversity of speculative sampling, we conduct a comparative analysis against Eagle2 (with temperature=1). Using a prompt from mt_bench: "Compose an engaging travel blog post about a recent trip to Hawaii, highlighting cultural experiences and must-see attractions." We statistically evaluate the responses averaged over 3 trials. The figure (https://anonymous.4open.science/r/Jakiro_rebuttal-C4C2/D.png) visualizes the frequency heatmap of the top-10 drafted tokens (default setting) sampled during the entire response generation. Tokens with frequencies below 10 were filtered out for clearer visualization.

The results show:

• Higher Token Quantity: Under identical experimental conditions, Jakiro generates significantly more draft tokens (vertical axis of the figure).
• Broader Semantic Coverage: The generated tokens occupy a wider span in the semantic embedding space (horizontal axis), reflecting richer topical diversity.

Additionally, we also provide the statistical results of accepted tokens for Jakiro* and Eagle2 on the MT-bench dataset. See https://anonymous.4open.science/r/Jakiro_rebuttal-C4C2/C2.png, which further highlights the diversity of our Jakiro.

W2: No efficiency or memory overhead analysis compared to non-MoE speculative methods.

We sincerely appreciate this insightful question. Here is our Memory Overhead Analysis:

Table: Hardware Metrics on A40-45GB of Vicuna 7B (measured by nvidia-smi as shown in https://anonymous.4open.science/r/Jakiro_rebuttal-C4C2/M.png)

Since Jakiro employs a lighter-weight MLP than Eagle, with merely 0.6% additional memory usage, the MoE-2 variant delivers a 10.6% speedup in token generation latency, showing highly efficient computation-memory scaling.

W3: The justification of Jakiro's design regarding its novelty and meaningfulness.

Jakiro's novelty lies not in simply applying MoE to speculative decoding but in developing a dynamic decoupling framework that fundamentally addresses a newly identified bottleneck. Here are key justifications:

• Problem Innovation
  Prior works focus on temporal decoupling (Eagle) or multi-head prediction (Medusa) but overlook in-step candidate correlation. This intrinsic limitation motivates our MoE-based decoupling at the intra-step level.

• Architectural Novelty
  Compared to standard MoE applications:

  • Dynamic Routing: Experts specialize in two-branch token speculative decoding.
  • Semi-autoregressive: Combines autoregressive decoding for early tokens and parallel decoding for later stages.
  • Contrastive MoE: The first to apply a contrastive mechanism between activated experts.

S1: The paper lacks an overview figure for the proposed framework.

Thanks for your suggestion. We will include the framework diagram (https://anonymous.4open.science/r/Jakiro_rebuttal-C4C2/F.png) in the revised version.

Q1: The relationship between the particular method in (a) part of Figure 1 and the other three methods.

Figure 1 (a) presents the baseline method of classical speculative decoding (i.e., SpS in Table 1 of this paper), which serves as the comparative reference for the improved approaches in (b) Medusa, (c) Eagle, and (d) Jakiro. We will explicitly annotate this relationship in the revised version.

Q2: What does Jakiro's name mean?

The name 'Jakiro' is inspired by the twin-headed dragon character from the DOTA game, symbolizing our method's dual-expert architecture where two specialized activated heads collaboratively generate diverse token predictions.