Thank you for your careful review of our paper and your insightful and constructive comments. We have addressed your comments and updated our manuscript accordingly. Please find our detailed answers below.

W1. Additional baseline - Missing Baseline - One-size-Fits-All (OFA) Drafter

Thank you for raising this important point. We conducted additional experiments with an One-size-Fits-All (OFA) drafter for both black-box and white-box settings. These OFA drafters were trained on the mixed datasets spanning all tasks, ensuring a direct comparison to the specialized-drafter-based MetaSD framework. The results in Table3 and Table4 of our revised manuscript show that while the OFA drafter performs well, our MetaSD framework outperforms OFA drafter across most tasks, demonstrating the strength of our method. The OFA Eagle drafter performs relatively well. However, it is still outperformed by MetaSD (i.e., MetaEagle-UCB).

Q1. Explanation on the advantage of MetaSD in batched inference setting

We acknowledge that the current statement in Appendix C.1 was not fully clear regarding our position on batched inference. We appreciate the opportunity to clarify our intention.

Our proposed MetaSD framework can address scenarios where tasks of varying nature (e.g., translation, summarization, QA, etc.) are mixed within a single batch. In such settings, a static approach that relies on a single drafter for all instances in the batch could lead to significant performance degradation, especially when the drafter is poorly suited to some tasks in the batch. This challenge is exacerbated as batch sizes increase.

MetaSD, in contrast, adapts dynamically at the instance level within a batch, ensuring that the most suitable drafter is chosen for each task. By leveraging our MAB framework for dynamic adaptation, MetaSD optimizes performance for diverse tasks within the batch, providing consistent improvements in throughput compared to static, single-drafter-based SD approaches. Thus, even as batch sizes increase, MetaSD maintains high throughput by avoiding the pitfalls of static, one-size-fits-all drafters.

We will revise Appendix C.1 to better articulate this point, as follows:

"Batched inference: Our current implementation primarily focuses on single-query scenarios. However, adapting the MetaSD framework to batched inference—where different tasks are mixed within a single batch—presents an opportunity for significant efficiency gains. Unlike static single-drafter-based SD, which can suffer from suboptimal performance when handling diverse tasks in a batch, MetaSD dynamically optimizes drafter selection at the instance level. This ensures consistently high throughput, even in high-throughput batched settings."

For more information, we would like to share the explanations regarding throughput efficiency in Appendix F.10. Our drafter management mechanism does not lead to throughput depletion compared to single-drafter methods due to the following reasons:

• No Increase in Memory Bandwidth Requirements: The drafters’ parameters are preloaded into DRAM and do not require frequent movement to VRAM for computation. This ensures that the number of memory movements and memory size for movement remain identical to single-drafter methods, even in scenarios with multiple drafters.

• Scaling Across Batches: Since the computational structure remains unchanged, performance observed in single-batch scenarios translates directly to multi-batch settings, maintaining throughput consistency.

We further conducted experiments comparing throughput following the same settings in the original Eagle paper. The results confirm that the performance of MetaSD scales well without significant degradation:

We hope this clarifies our approach and the rationale behind our claims. Thank you for highlighting this important aspect!

Thank you for your question. To clarify, the throughput experiments were conducted on a mixed dataset comprising tasks from diverse domains, including Code Generation, Math, Question Answering (QA), Translation, and Summarization.

The throughput measurement methodology strictly follows the settings described in the original Eagle paper. This was done intentionally to enable direct comparison, highlighting the improvements achieved by our MetaSD framework over existing approaches.

Additionally, we would like to emphasize that our method does not critically impact the roofline model or memory bandwidth (as we mentioned above). MetaSD efficiently manages memory usage, ensuring that there are no significant increases in memory movement or bandwidth demands. This efficiency contributes to the strong throughput performance observed in our results, and we believe this further supports the validity of our framework.

We hope this explanation addresses your concerns and provides additional context for the results. Please let us know if you have further questions!

-Authors

Follow-up. Throughput

By the way, before discussions, we apologize to Reviewer CxuG for the extended discussion with Reviewer V7AR in this thread and appreciate your understanding.

Thank you for pointing this out. We would like to clarify that the throughput experiments reported in our revised manuscript were conducted under homogeneous batch settings, where all tasks within a batch were identical, while the reported value is the average performance across diverse tasks.

To address your concerns, we conducted additional experiments to assess the performance of MetaSD in the worst-case edge scenario, where each instance in the batch corresponds to a different task (heterogeneous setting). Below are the results:

Observations

1. Homogeneous batches: In scenarios where all instances in the batch belong to the same task, MetaSD achieves consistently higher throughput than the OFA Eagle. This is due to the ability of MetaSD to dynamically select the best drafter for the task, ensuring optimal performance.

2. Heterogeneous batches (Worst-Case Edge): In the edge case where each instance in the batch corresponds to a different task, we observe a slight drop in throughput for MetaSD. This is expected due to increased I/O from switching between drafters. Notably, for larger batch sizes, the performance of MetaSD in this setting may fall below the OFA Eagle, highlighting a limitation of our approach in such cases.

3. Small batch sizes: For small batch sizes (e.g., batch size = 1 or 2), MetaSD remains highly compatible even in heterogeneous task settings.

Discussion

We acknowledge that the performance drop in heterogeneous batch scenarios represents a limitation of our approach. However, this is one of edge cases that may not frequently occur in real-world applications. Moreover, this is not a problem unique to MetaSD—other systems relying on specialized drafters may face similar challenges in such settings.

On the other hand, MetaSD demonstrates strong adaptability and efficiency in mixed-task, small-batch scenarios, which are also common in many practical deployments (e.g., laptop, personalized products). Its ability to dynamically select the most suitable drafter ensures robust performance, even in uncertain or evolving task distributions.

Broader Context and Future Directions

• One of the key motivations behind MetaSD is its applicability to real-world scenarios where multiple heterogeneous drafters are available, but their training data or performance characteristics are unclear. In such cases, our approach provides a safe and reliable choice, dynamically selecting the best drafter based on real-time feedback.

• Additionally, as noted in our paper, the framework can potentially be extended to higher-level optimization tasks. For instance, instead of focusing solely on drafter selection, the framework could be adapted to dynamically choose the optimal speculative decoding algorithm itself. We believe this versatility represents a promising direction for future work.

In the camera-ready version, we will include these findings and further clarify the advantages and limitations of MetaSD under different batch scenarios. Thank you for highlighting this aspect, which allowed us to explore and communicate the broader implications of our framework.

Let us know if there are further questions or areas where we can provide additional clarity.

Thank you for your careful review of our paper and your insightful and constructive comments. We have addressed your comments and updated our manuscript accordingly. Please find our detailed answers below.

W1. Extra computation overhead of MetaSD

We would like to address each point in detail to clarify potential misunderstandings regarding the computational implications of our approach.

Training Overhead

While specialization may require additional training efforts compared to an OFA (One-size-Fits-All) drafter, we would like to emphasize that our approach can handle real-world scenarios where heterogeneous drafters already exist in public repositories. Our framework focuses on optimizing the utilization of such heterogeneous drafters, ensuring that the most suitable drafter is selected dynamically. This shifts the problem from retraining models to developing an effective strategy for utilizing pre-existing resources. Thus, while training specialized drafters may involve additional costs in some cases, the broader applicability and versatility of MetaSD provide substantial practical value. In addition, issues regarding the cost of training drafters are not solely limited to our work, but all research regarding this area.

Inference Memory-Bandwidth Efficiency

We would like to clarify a potential misunderstanding regarding inference memory bandwidth requirements. Although MetaSD employs multiple drafters, this does not increase the memory bandwidth requirements during inference. The drafters’ weights are preloaded into DRAM, ensuring no additional VRAM bandwidth is consumed compared to a single-drafter setup. Memory bandwidth refers to the data movement between VRAM and compute cores, which remains identical regardless of the number of drafters, as only one drafter operates on the VRAM-resident data at a time. The following table provides a comparison:

By ensuring that only the active drafter interacts with the VRAM, MetaSD does not increase VRAM bandwidth demands, maintaining parity with single-drafter approaches during inference.

Serving Complexity

Using multiple drafters in MetaSD does not inherently increase serving complexity. Modern distributed systems already employ model parallelism techniques to allocate workloads effectively across multiple GPUs. In MetaSD, the drafters are evenly distributed across available GPUs, with each GPU independently handling its assigned drafter without added coordination costs. This design ensures that:

• Load Balancing: Drafters are distributed across GPUs based on their assigned tasks, maintaining equivalent complexity to single-drafter systems.

• Minimal Communication Overhead: MetaSD requires no additional inter-GPU communication beyond standard model parallelism setups.

Consequently, the serving complexity of MetaSD aligns with that of traditional parallelized single-drafter systems.

Justification of Overhead

The modest increase in DRAM memory usage (+2 GB) and marginal training cost for specialized drafters is justified by the significant performance gains achieved through adaptive optimization. MetaSD dynamically selects the most suitable drafter for each task, consistently outperforming single-drafter methods across a wide range of scenarios, as highlighted in our experimental results.

Furthermore, MetaSD's adaptability addresses an important real-world challenge: utilizing publicly available, pre-trained heterogeneous drafters effectively. By offering a generalizable strategy for optimizing these resources, MetaSD provides practical value beyond specialized retraining, supporting diverse and evolving task requirements.

Q1. Clarification on computing KV-caches

In our paper, the KV cache is recalculated for the previous context whenever a drafter switch occurs. However, we found that this recalculation introduces negligible computational overhead, even for relatively long contexts. This is due to the limited cost associated with prefilling the KV cache of a small drafter (e.g., In the case of Eagle drafter, one layer of KV cache is computed for unseen context), which is highly efficient.

Additional Experiments for Efficient KV Cache Strategies

To further validate our framework's efficiency, we additionally experimented with incorporating StreamingLLM [2], which can avoid full KV cache recalculation. The results demonstrate that this approach achieves comparable performance, confirming the robustness of MetaSD:

Key Observations

1. Minimal Impact of KV Recalculation: The results show that even with full KV cache recalculation, MetaEagle-UCB maintains strong performance, highlighting that this process introduces minimal computational overhead.

2. StreamingLLM Effectiveness: Incorporating StreamingLLM techniques marginally improves performance in some cases (e.g., Translation and QA), while maintaining a similar overall performance profile.

Q2. Additional baseline

We conducted additional experiments with an One-size-Fits-All (OFA) drafter for both black-box and white-box settings. These OFA drafters were trained on the mixed datasets spanning all tasks, ensuring a direct comparison to the specialized-drafter-based MetaSD framework. The results in Table3 and Table4 of our revised manuscript show that while the OFA drafter performs well, our MetaSD framework outperforms OFA drafter across most tasks, demonstrating the strength of our method. The OFA Eagle drafter performs relatively well. However, it is still outperformed by MetaSD (i.e., MetaEagle-UCB).

References

[1] Leviathan et.al., Fast Inference from Transformers via Speculative Decoding. ICML 2023.

[2] EFFICIENT STREAMING LANGUAGE MODELS WITH ATTENTION SINKS. ICLR 2024.

Thank you for your detailed response!

It has addressed most of my concerns, and I do appreciate the huge efforts you’ve put into this paper and the response. Therefore, I'm raising my score to 6. That said, I have some follow-up questions and suggestions for further improvement:

• Comparisons with Concurrent Work: I recommend including discussions of related concurrent work, such as SWIFT[1]. SWIFT dynamically optimizes and switches to a more efficient drafter using layer skipping, which shares similarities with the Bandit approach in this manuscript. Additionally, I encourage evaluating the Bandit approach in a dynamic data stream with multi-task OOD data, as demonstrated by SWIFT in Figure 7. This would provide a more comprehensive evaluation (you don't have to do it during the rebuttal period). Besides, more discussions with [2] should be included since this method also optimizes SD performance on the fly.
• Computation Overhead: The analysis of additional computational overhead introduced by the Bandit mechanism should be included in the revised manuscript to provide a clear understanding of its practical implications.
• Regarding W2: I respectfully disagree with the statement, "Computing similarity between sentence embeddings requires passing the context through an encoder-only network to obtain the embeddings." Cosine similarity can be directly calculated on the input representations after the pre-filling stage of the target LLM, without requiring an additional encoder. However, I acknowledge the authors’ valid point that this naive strategy could pose challenges for tasks with highly similar representations.
• Regarding W4: The manuscript should report switching accuracy metrics. Specifically, I am curious whether the low switching frequency observed leads to high switching accuracy (i.e., the Bandit mechanism consistently selects the optimal drafter). If this is not the case, the low switching frequency could undermine the effectiveness of the Bandit mechanism, and further clarification or analysis would be valuable.

Thank you again for your thorough response. I look forward to further discussions and may consider increasing my score upon addressing these points.

[1] SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration. Xia et al. Arxiv 2024.

[2] Online speculative decoding. Liu et al. ICML 2024.

Thank you for your thorough and detailed responses. I have no further concerns and have updated my score accordingly.

Thank you for your careful review of our paper and your insightful and constructive comments. We have addressed your comments and updated our manuscript accordingly. Please find our detailed answers below.

W1. Invalid evaluations for real-world use cases.

We acknowledge the concern that using the same prompt template across all instances within a dataset might not fully reflect real-world scenarios. To address this, we conducted additional experiments with perturbed prompts, where the prompt for each query varied slightly but was semantically equivalent to the original. These perturbed prompts were generated using GPT-4, ensuring natural and diverse variations. For example:

• Translation Task: Original prompt “Translate German to English” was perturbed to “Convert this text from German to English.”

• Summarization Task: Original prompt “Summarize:” was perturbed to “Provide a concise overview of the following text:”

Key Observations

The main results and descriptions are detailed in Table 10 and Table 11 of Appendix F.9.

1. Performance Impact Across All Methods:

   • Perturbed prompts led to a drop in performance for all drafters, including One size Fits All (OFA; which is trained on a mixed dataset) and individual specialized drafters. This indicates that the issue is not limited to our approach but affects all SD methods.
   • The observed degradation reinforces that real-world variability in prompts can challenge any static drafter selection strategy.

2. MetaSD’s Robustness:

   • Despite the variability introduced by perturbed prompts, MetaSD consistently outperforms all baselines, including OFA and individual drafters. This result highlights the advantage of MetaSD’s dynamic selection mechanism, which adapts to token distributions at every round rather than relying solely on prompt characteristics.

W2. Limited effectiveness of the MAB approach

While the datasets used in our experiments are indeed distinctly different and well-defined, our method is designed to handle more complex, real-world scenarios where:

• The boundaries between datasets or task types are less clear.

• The performance of each drafter can vary dynamically depending on the token distribution as the decoding progresses.

• Out-of-domain tasks introduce additional complexity, which static rule-based systems or simpler machine learning models cannot handle effectively.

While it is true that MetaSD involves SD rounds for multiple drafters, the additional computational cost is mitigated by the efficiency gains in overall throughput. The reported accuracy of drafter selection (below 80%) should not be viewed in isolation, as it reflects the complexity of real-world token prediction rather than a static task classification problem. While rule-based systems may appear efficient in well-defined scenarios, they lack the flexibility to adapt to:

• Mixed-domain or evolving tasks where token distributions do not align clearly with predefined rules.

• Real-time feedback from the decoding process, which is a critical component of MetaSD's adaptive optimization.

Thank you for elaborating on your concerns. We would like to clarify and address the points raised.

1. On Real-World Use Cases and Prompt Variations

We want to emphasize that the experiments with perturbed prompts were conducted by varying the prompts for every query, not just using a single variation. The example provided in the response was one of many variations used. The results reflect this setup and demonstrate that the performance degradation affects all methods, not just MetaSD. This highlights a broader challenge that is not unique to our approach (Most methods face this challenge).

2. On the Practical Value of MetaSD

While it is true that MetaSD-UCB does not outperform individual specialized draft models in every case, its strength lies in its robustness across diverse scenarios, including in-domain, out-of-domain, and perturbed settings. These scenarios reflect the complexity of real-world tasks, where heterogeneous drafters are often pre-trained on unknown data or lack clear task-specific boundaries.

Although MetaSD does not consistently outperform specialized draft models on their respective tasks, it is highly robust across diverse tasks. On average, MetaSD achieves performance comparable to specialized models when evaluated across a mix of tasks. For context, we analyzed MetaSD's results by comparing them to a simulated classification-based approach, where the classifier is assumed to have zero inference time, and found that MetaSD effectively achieves:

• 87% of the classification accuracy of specialized drafters in MetaSps settings.
• 90% of the classification accuracy in MetaEagle settings.

This demonstrates that MetaSD delivers near-specialized performance from the perspective of classification, without relying on pre-trained parameters (already working as a good classification predictor). Even without offline training or fine-tuning, MetaSD consistently provides strong, adaptive performance across diverse tasks. Moreover, this robustness extends to out-of-domain and perturbed scenarios, which often present significant challenges for classification-based or static systems. On the another line, MetaSD can be also supported by multiple theorems, which can be huge difference from the classification-based system.

Additionally, MetaSD’s reliance on the multi-armed bandit (MAB) algorithm—without any learnable parameters—makes it highly practical and computationally efficient for real-world on-the-fly applications.

3. On Classification-Based Approaches

We appreciate your suggestion regarding classification-based approaches and agree that they could serve as a useful baseline. However, as noted, classification models rely on pre-trained embeddings and are subject to classification errors. This introduces challenges, particularly in:

• Out-of-Domain Scenarios: Classification models often generalize poorly to tasks outside their training domain.
• Dynamic Adaptation: Classification requires offline pre-training and cannot adapt to real-time token-level feedback (on-the-fly), unlike MAB approaches.

For camera-ready, we plan to evaluate a simple classification-based approach using Google-BERT-Base-Uncased embeddings with a linear layer to predict the optimal drafter for each query. This experiment will provide an additional baseline for comparison.

4. Broader Motivation and Contributions

As highlighted in our revised manuscript, one of the key motivations of MetaSD is to address scenarios involving heterogeneous drafters where training data and performance characteristics are unknown. MetaSD provides a robust and simple alternative for on-the-fly decision-making, avoiding the need for pre-trained models or extensive offline fine-tuning.

While classification-based methods may serve as a useful baseline, their inability to dynamically adapt makes MetaSD a practical and effective solution, particularly in scenarios requiring robust performance across diverse tasks.

Dear Reviewer Tyx3,

In response to your concerns regarding classification-based routing, we conducted additional experiments using BERT-based models for routing across diverse tasks (coding, translation, summarization, QA, and math reasoning). Specifically, we evaluated routing performance using pre-trained BERT-Base-Uncased (110M) and BERT-Large-Uncased (330M) with both cosine similarity and fine-tuned softmax heads. All experiments are conducted with a single A100 GPU. (Note that by considering the small drafter used in SpS has 68M parameters, both models are quite large.) The results are as follows:

• Cosine Similarity (Pre-trained):
  • BERT-Base: 24.50% accuracy
  • BERT-Large: 30.00% accuracy
• Fine-Tuned BERT with new Softmax Head:
  • BERT-Base: 74.54% accuracy
  • BERT-Large: 85.58% accuracy

We also measured the average latency speedup across all tasks:

Average Latency Speedup Ratio (In-domain)

Our MetaSD demonstrates higher robustness than classification-based routing, particularly in perturbed and out-of-domain (OOD) scenarios.

Perturbed Prompts Speedup Ratio

Out-of-Domain Speedup Ratio (RAG and Finance)

These results demonstrate that our framework is more robust and adaptable, particularly under challenging conditions like perturbed prompts and OOD tasks. Classification-based routing struggled with task confusion, especially between closely related tasks. For example, coding and math often share overlapping patterns in token distributions, making them harder to distinguish, while summarization and QA frequently confused sentence-level context. Even though translation appeared to perform well in classification, it was likely aided by distinct vocabulary distributions unique to translation tasks, which reduce ambiguity. (Note: While classification might be better suited for multilingual tasks due to distinct vocabularies, in other cases, it was not as effective.)

A key reason why MAB demonstrates stronger performance in this setting lies in our novel block-divergence reward. This reward effectively captures the distributional alignment between the target model and the drafter, allowing the MAB to dynamically adapt and select the optimal drafter based on real-time feedback (i.e. policy-routing). In contrast, classification-based routing does not consider this alignment, leading to suboptimal routing decisions that fail to exploit token-level distribution patterns.

Of course, we believe that classification-based models could potentially address these issues more effectively if designed with a deeper consideration of token-level distributional differences. However, our focus is not on developing classification-based routing methods, and to the best of our knowledge, no established approaches currently exist in this space. Instead, we proposed a simple but efficient method for the novel problem (that we newly raise in this field) at hand. We also believe that exploring classification-based routing as a future direction could offer an orthogonal and promising approach to improving routing efficiency.

We hope this detailed analysis clarifies the advantages of our approach and demonstrates why a simple classification-based method is insufficient. We will work to include these in the camera-ready.

Thank you for your time and consideration.

Best regards,

-Authors

Nonstationary Environment

To further demonstrate the effectiveness of the MAB approach, we conducted experiments on a non-stationary translation task, where each query required translating two different languages (French and Chinese) into English. The experiments were performed using a single RTX 3090 GPU, and the results are as follows:

These results clearly demonstrate the effectiveness of the MAB approach, where MetaSD-UCB with $\alpha=0.1$ achieves a speedup ratio of 1.722, which even exceeds the performance of the optimal drafter (1.581). This improvement arises because MAB dynamically adapts to the best drafter in the current environment through a combination of exploration and exploitation.

For example, in a multilingual scenario, where the task is "Translate French and German to English, respectively," classification-based routing would be limited by the performance of the best single specialization (i.e., the speedup ratio of the optimal drafter as its upper bound). In contrast, MetaSD-UCB surpasses this upper bound by effectively adapting its policy dynamically, demonstrating a unique advantage over classification-based methods in non-stationary environments.

Future Directions

We are already encouraging further exploration in our paper's Future Work section, and we believe this type of policy-routing could benefit significantly from more advanced approaches like contextual bandits or reinforcement learning (RL). These methodologies may offer more sophisticated mechanisms for dynamic adaptation in non-stationary environments. Furthermore, from a scalability perspective, interesting future directions include:

• Exploring how to cluster drafters effectively for system-level applications, such as those used in retrieval-augmented generation (RAG).
• Investigating how to design OFA drafters in an orthogonal manner to maximize complementarity across tasks.

Clarification on Our Research Focus

It’s important to highlight that our work does not advocate training multiple specialized drafters instead of OFA drafters. One of our main focus is to address the critical challenge of dynamically routing tasks among pre-existing heterogeneous specialized drafters. This is a practical and timely problem, especially as these specialized drafters are becoming widely available.

For a deeper understanding of our motivation and scope, we strongly recommend revisiting the Motivation Section in the main paper and the Further Motivation Section in the Appendix. These sections provide a comprehensive view of why this research is both relevant and necessary for advancing speculative decoding systems.

Thank you for your careful review of our paper and your insightful and constructive comments. We have addressed your comments and updated our manuscript accordingly. Please find our detailed answers below.

W1. Clarity of the assumption on acceptance rate

We denote the acceptance rate of drafter i in round t as $\alpha_{i,t}$ and assume this is a stationary random variable with mean $\alpha_i$. As the regret bound of a bandit algorithm is usually related to the expectation of each reward, $\alpha_i$ appears in the result in Theorem 2 but this does not mean we replace $\alpha_{i,t}$ with $\alpha_i$ in any part of our analysis. For clarification, we add a formal definition of our assumption on acceptance rate in Appendix G.6.

W2. Validity of Theorem 1

The expectation of BD reward in Lemma 4 is a direct consequence of our assumption. Moreover, our stationarity assumption implies that the TV distance between target model and i-th drafter does not depend on the context and consequently all of the theorem does not have to take into account the target model output. Please note that our assumption is stricter than the most general case as in [1] while it is more general than [2] where authors use fixed acceptance rates. For completeness, we add the comparison between ours and assumptions used in [1] and [2] with analysis in Appendix G.6.

W3. Stochasticity of the target sequence length B

Thank you for the constructive feedback! While we define the stopping time regret (Definition 2) and build Theorem 2 with fixed B for each generation, this can be easily generalized to the more general case considering a probability space induced by the target model. This is possible since target sequence length B does not depend on the policy under our assumption where target sequence length B is only determined by the randomness of the target model, independent from the types of drafters we used throughout the generation. We include a general version of Theorem 2 with analysis in Appendix G.7.

W4. The relationship between the reward / regret design and the performance of the algorithm.

The performance of a bandit algorithm itself does not depend on the definition of the regret since regret objective is only for the analysis. We show that the optimality guarantee with the existing regret objective in the bandit literature will not prove the optimality of the algorithm performance in the Speculative Decoding. This motivates us to design a novel regret objective in Definition 2 which aligns with the performance of algorithms in SD. This is reflected in Lemma 6, where we prove the upper bound on existing regret definition may not guarantee the optimal performance. One of our key theoretical contributions is defining this new regret objective and proving the upper bound on this new stopping time regret which is not so straightforward from the conventional analysis of UCB. Consequently, the result of Theorem 2 reflects the guarantee of the algorithm performance and lower regret upper bound when using BD reward can imply better algorithm performance in most scenarios which we state in Corollary 1.

W5. Line 1586 and Line 1695

We appreciate your careful review. We fixed typos in line 1763 (originally line 1586) and line 1878 (originally line 1696) of our original manuscript.

For Line 1586, $\mathbb{E}[\tau^u(\pi^u,B)]$ should be changed to the $\mathbb{E}[n_{i^{\star}}(\pi^u,B)]$.

For Line 1878 (eq. 29), the square root term in the left hand side should be divided into $s$ and $n_i$, respectively. As in the proof of the original UCB algorithm (proof of Lemma 7), confidence bound in both inequalities needs to be shrinked as $s$ and $n_i$ increases with fixed t and it doesn’t necessarily have to be shrinked with increasing t.

W6. Hyperparameter $\beta$ and analysis on mean gaps $\Delta_i$

Thank you for the insightful comments!

In Appendix G.8, we include further analysis on how the regret upper bound of our algorithm changes with different $\beta$.

Also, mean gaps $\Delta_i$ can indeed be a critical factor of the algorithm performance. We make terms containing mean gaps to appear in inequality in Theorem 2 and relevant equations to better represent this fact which ensures all of the constant terms to be independent both on $B$ and mean gaps.

W7. Formal definition of B and terminology

We provide a formal definition of target sequence length B in Appendix G.7. We have also ensured consistent use of the term ‘target sequence length’ when referring to $B$ throughout the manuscript to avoid any ambiguity.

Q1. The performance of EXP3

While UCB algorithm is optimal in instance dependent regret bound with stationary assumption, EXP3 algorithm explores more on suboptimal arms. We also observed using EXP3 algorithm results in picking suboptimal drafters more compared to using UCB algorithm as our ablation study on the best arm ratio (Section 4.3) shows.

Q2. Additional baseline for mixed dataset

We conducted additional experiments with an One-size-Fits-All (OFA) drafter for both black-box and white-box settings. These OFA drafters were trained on the mixed datasets spanning all tasks, ensuring a direct comparison to the specialized-drafter-based MetaSD framework. The results in Table3 and Table4 of our revised manuscript show that while the OFA drafter performs well, our MetaSD framework outperforms OFA drafter across most tasks, demonstrating the strength of our method. The OFA Eagle drafter performs relatively well. However, it is still outperformed by MetaSD (i.e., MetaEagle-UCB).

Q3. Scaling law of the task numbers of MetaSD

While our experiments focus on five tasks, we selected them to cover orthogonal areas (e.g., code, translation, QA, summarization, and math) to ensure diverse task representation. We acknowledge that five tasks may not fully capture the broader task spectrum. However, this limitation is not unique to our framework but also affects single-drafter approaches. To answer your request, we conducted out-of-domain experiments with Alpaca-Finance and RAG datasets in Table 9 of Appendix F.8. Table 9 demonstrates that even with only five specialized drafters, MetaSD consistently outperforms single-drafter approaches, including OFA drafters, under this scenario.

Reference

[1] M. Yin et al, A Theoretical Perspective for Speculative Decoding Algorithm. NeurIPS 2024.

[2] Leviathan et.al., Fast Inference from Transformers via Speculative Decoding. ICML 2023.

W6. Comparison with ETC

We appreciate you for the constructive feedback with great intuition.

Indeed for a small $\beta$, a regret upper bound including $\beta$ does not hold anymore. We conducted additional experiments on translational task using ETC with tuned exploration rounds as the reviewer’s suggestion and the result is as follows:

In the above result, ETC(15) is based on the uniform exploration round of $15$ which we found performs best empirically. And exploration hyperparameter $\beta$ is $0.1, 0.01$ are used for a fair comparison. While ETC achieves comparable speed-up performance compared to UCB, ETC requires the number of exploration rounds as a hyperparameter and the optimal exploration round depends on the total target sequence length $B$ which we can’t know during the generation. We will include the discussion with updating the result into our final manuscript to include ETC as a baseline of our experiments.

W7. Clarification on the statement

Thank you for pointing out this. As stated in the additional response of W3, in order to relax the assumption on fixed B and make Theorem 6 to be true, we have to rely on the additional assumption (Assumption 2 in Appendix G.7). For clarification, we restate the previous claim into "We can consider general scenarios where we take all possible instances generated by a target model when using temperature sampling with T>0. In this scenario, we define the expected regret over the probability space induced by the target model." in our revised manuscript.

Additional comparison with [1] and clarification

We appreciate the reviewer for providing us good reference [1] which we didn’t have a chance to access by the time of our submission. We believe analyzing Speculative Decoding on the more general assumption of formalizing a decoding problem using a Markov-chain would be closer to the real-world scenarios. While formulating our problem as a non-stationary bandit with Markov chain formulation as in [1] would be a more realistic modeling, non-stationary bandit does not always perform better empirically.

Here, we would like to emphasize once more that our work is the first to investigate multi-drafter selection within one generation when doing Speculative Decoding and we especially take advantage of a bandit algorithm which is simple yet effective. Thus future works are expected to investigate more general scenarios and corresponding analysis to cover more realistic modeling.

We sincerely appreciate you for the constructive discussion and valuable feedback with excellent intuitions. We hope our responses clarify your initial concern and we are open to further discussion if there exist any remaining concerns or questions!

Reference

[1] M. Yin et al, A Theoretical Perspective for Speculative Decoding Algorithm. NeurIPS 2024.

[2] Leviathan et.al., Fast Inference from Transformers via Speculative Decoding. ICML 2023.

Thank you once again for your thorough review. In response to your follow-up comments, we have further revised and updated our manuscript accordingly.

W1. Explicit statement of i.i.d. assumption

Thank you for the suggestion, in revised manuscript, we explicitly state that $\alpha_{i,t}$ are i.i.d from a distribution $\nu_i$ with mean $\alpha_i$.

W2. Explicit statement of i.i.d. assumption in Theorem

In the revised manuscript, we also directly state that $\alpha_{i,t}$ are i.i.d from a stationary distribution $\nu_i$ in Theorem 2 to properly reflect on our i.i.d. assumption. Moreover, we will include a comparison for Assumptions in [1] and [2] in our main paragraph (Section 3) for further clearness.

W3. Stochasticity of the target sequence length B

We appreciate the constructive feedback and further clarification.

First, we would like to clarify that Assumption 1 is about assuming acceptance rates $\alpha_{i,t}$ are i.i.d. from the same distribution $\nu_i$ for each instance of generation. While this distribution $\nu_i$ can vary across generations and depend on the target sequence length B, this does not impact our algorithm's performance under Assumption 1 because we reset the bandit instance for each new generation. Suppose a target mode (with the same rule in your example) generates 1110. Then the acceptance rate of a drafter $i$ conditioned on 1, 11, 111, 1110 follows from the same distribution $\nu_i$ with mean $\alpha_i$. Thus, the objective of a bandit algorithm here is to find the optimal drafter and minimize the regret within this instance and our analysis holds in this case. Next, suppose the target model generates 110 then EOS. Then, acceptance rates of drafter $i$ for this instance also follows stationary distribution with mean $\alpha_{i'}$ which can be different from $\alpha_i$, which does not affect our algorithm since the policy in the instance is independent of any previous results. In this sense, Theorem 2 is not based on the assumption of fixed $B$ in generation, rather, the regret is defined for each instance of generation. Consequently, our analysis of Theorem 2 remains valid under Assumption 1.

However, to further enrich our results, we take the reviewer’s advice and include Theorem 6 for the general analysis on the expected regret where expectation is taken over the probability space induced by the target model. This should be built on an additional assumption about assuming $\alpha_i$ being independent of B and its conditional expectation over a given $B$ being the same for every $B$. We explicitly state this assumption before Theorem 6 to avoid any confusion with Assumption 1 and Theorem 2.

For the last part, we correct our statement to "Since drafter selection from the policy $\pi$ is independent from $B$ under Assumption 2".

W4. Temperature sampling

Thank you for raising the important point.

Assumption 1 can hold with every $T$ under our formulation and the analysis holds in Greedy Decoding with Assumption 1. With this, our theoretical results include the case of $T=0$, greedy decoding. We inlcude this in our manuscript for avoid any confusion.

The assumption is based on our empirical observation that TV distance between two probability distributions can be an intrinsic measure of closeness of the target model and a drafter and assumes acceptance rates of a drafter are i.i.d. drawn from the same distribution in each instance. This simplification allows us to investigate theoretical guarantee and reward design when using a bandit algorithm in multi-draft SD.

Moreover, in addition to Table 7 where we showed that our algorithm still performs strongly with $T>0$, we will include further experimentally results from the sampling decoding in all of our main experiments in our final manuscript to connect the theory and the experiments more robustly.

Dear Reviewer v7AR

Thank you for your suggestion. We will do our best to conduct as many additional experiments as possible within the discussion timeline, including out-of-domain datasets like MMLU, Physics QA, and Hotpot QA, with batch size > 1 and OFA Eagle comparisons.

That said, we feel this level of experimental demand is quite extensive compared to typical expectations, as similar works, even including concurrent studies at ICLR, are not usually evaluated across such a broad range of datasets.

Additionally, we would like to ask whether the lack of specific comments on the theorem and its implications is a critical factor in your decision to maintain the current score (3: reject, not good enough). While we believe our responses (and revised manuscript) address the concerns raised, given the tight timeline, we want to ensure we prioritize addressing the most impactful concerns if others remain.

Warm regards,

-Authors

Dear Reviewer v7AR,

Thank you for your insightful questions. Below, we address your concerns about the relationship between the theory and experiments.

We have conducted additional experiments with stochastic decoding (temperature=1.0), as suggested, to demonstrate the robustness of our framework under the setting assumed in the theoretical analysis. The results are as follows:

Black-Box (MetaSpS)

White-Box (MetaEagle)

These results align with the theoretical assumptions, still demonstrating the robustness of MetaSD under stochastic decoding conditions. Additionally, Table 7 in our manuscript already presents the robustness of MetaSD in multilingual settings under stochastic decoding. We believe these results further clarify the connection between our theory and experiments.

To ensure clarity, we will expand on these findings in our camera-ready manuscript and address any potential ambiguities that may arise for other readers.

Thank you for your follow-up comment. Respectfully, we believe there is a fundamental misunderstanding regarding the focus and intent of our work.

Respectfully Addressing the Concerns

1. Research Focus

Our work does not advocate for training multiple specialized drafters instead of OFA drafters. Rather, it addresses a critical and practical challenge: how to effectively utilize pre-existing heterogeneous specialized drafters dynamically. This is a practical scenario as specialized drafters, trained on diverse datasets, are becoming increasingly available. Routing tasks dynamically among these drafters is the focus of our research.
We strongly recommend revisiting the Motivation Section and Further Motivation Section in the Appendix, where we have explicitly detailed this objective, please.

2. On MAB’s Complexity

We respectfully disagree with the notion that MAB is a complex algorithm. MAB is one of the simplest and most traditional online learning approaches, widely used in practical domains like recommendation systems, finance, and medical systems for its adaptability and low computational overhead. These strengths make it an ideal choice for speculative decoding in dynamic environments.

3. Role of OFA in Routing

While OFA drafters are highly versatile, they are not guaranteed to perform optimally across all tasks. Notably, in our framework, OFA itself can be a key component in the routing process, serving as a fallback drafter to handle tasks where specialized drafters may perform suboptimally. This ensures robustness in scenarios where potentially low performance could arise. However, the intent of our paper is not to promote a single OFA drafter for all tasks, but to explore how pre-existing multiple (specialized) heterogeneous drafters [1], including OFA, can be used together dynamically to achieve better overall performance.
The comparison with OFA Eagle in our experiments serves to evaluate its effectiveness relative to MetaSD, not to suggest that OFA should be excluded. Rather, OFA is complementary and could play a vital role within a routing framework to ensure robustness.

• [1] Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters, Yi et al., EMNLP 2024.

4. Throughput and Real-World Applicability

We have addressed throughput concerns in our experiments. In single-batch scenarios, MetaSD-UCB consistently outperforms OFA Eagle, particularly in perturbed prompt settings. In heterogeneous batch throughput, MetaSD-UCB remains competitive due to its ability to utilize task-specific drafters dynamically. This adaptability is essential for handling real-world, diverse queries effectively.
Moreover, in multilingual translation or complex tasks like MMLU-COT and Physics QA, our framework dynamically leverages specialized drafters for specific tasks while complementing them with OFA as needed. Routing ensures that each task benefits from the most appropriate drafter, whether specialized or OFA.

• We respectfully request not to frame throughput as the sole determining factor for evaluating real-world applicability. In practice, single-batch inference and homogeneous batching are commonly used in industry applications, especially for scenarios like low-latency, user-specific generation.
• Additionally, if throughput is always the primary criterion, we would like to ask whether prior speculative decoding (spec-dec) papers—many of which are known to have lower throughput [2]—should also be discounted based on this metric alone. We believe such an approach overlooks the broader contributions and utility of speculative decoding frameworks, including improvements in latency, task adaptability, and system-level design.

[2] Fast Inference from Transformers via Speculative Decoding, Leviathan et al.

5. The Core Contribution

Our framework is a generic and scalable solution to the problem of dynamic task routing among heterogeneous drafters. It is not meant to replace OFA but to complement scenarios where pre-trained heterogeneous drafters exist—a situation that is increasingly common. This is a system-level problem that goes beyond a single-drafter approach and requires effective, adaptive solutions.

Final Thoughts

Respectfully, we believe the comments focus on a preference for OFA without fully engaging with the central problem our paper seeks to solve. OFA’s inherent limitations in non-stationary, multilingual, and domain-diverse settings underscore the importance of dynamic task routing. Furthermore, OFA itself can be effectively utilized within our framework to guard against potentially low performance, reinforcing the versatility of MetaSD.

We welcome further constructive discussion and hope this response clarifies the contributions and implications of our work.

-Authors

Thank you for the ongoing discussion. We would like to raise a critical concern:

• Request for Moderation: To the AC, we are concerned that certain interpretations appear to frame the work unfairly, not due to a lack of engagement but rather through persistent misdirection of the discussion. This includes steering rebuttal efforts toward edge cases and issues tangential to the core problem our work aims to address, while also demonstrating a misunderstanding of the theorems despite repeated clarifications. We kindly request your moderation to assess whether this critique appropriately reflects the context and intent of our contributions. If you find similar concerns, we sincerely ask for your guidance to ensure fairness.

We appreciate your attention to this matter as we strive for a balanced and constructive review process.

Best regards,

-Authors

Dear Reviewer v7AR,

We conducted additional evaluations regarding heterogeneous batch throughput and OOD datasets (adding to the previous finance and RAG OOD dataset evaluations), and the results are as follows:

Single Batch Throughput

Heterogeneous Batch Throughput

Analysis

MetaEagle-UCB consistently outperforms OFA Eagle in single-batch scenarios across Physics QA and Hotpot QA. MetaEagle-UCB performs better on average across the 57 tasks in MMLU-COT and outperforms OFA Eagle in 37 individual tasks while trailing in 20 tasks (Here, MMLU with CoT reasoning follows the setting by Anthropic). For heterogeneous batch throughput, MetaEagle-UCB remains competitive due to using its specialized drafters (e.g., for QA, code, and math; at most 3 drafters in general), which excel at handling diverse tasks within the MMLU dataset.

These findings suggest that future research could explore how to structure specialization among drafters to further enhance performance for throughput. For instance, incorporating two distinct types of OFA drafters trained on orthogonal domains could improve coverage across diverse task distributions, especially in mixed or OOD scenarios. This approach could also serve as a robust foundation for dynamically balancing specialization and generalization within speculative decoding frameworks.

While the current specialization approach shows slight limitations in some OOD scenarios, our research focuses on proposing a generic framework rather than fixed specialization. We believe that future studies, emphasizing how to optimally construct and organize domains during training, could effectively address these limitations and further improve performance in diverse and challenging conditions.

Clarifications on Research Focus

Our work is not advocating for training multiple specialized drafters instead of OFA drafters. The main focus is on effectively utilizing pre-existing heterogeneous specialized drafters, tackling the critical challenge of dynamically routing tasks among them. This is a practical problem, especially as such specialized drafters are becoming widely available.

• We strongly recommend revisiting the Motivation Section and Further Motivation Section in Appendix.

Broader Implications

This approach mirrors routing strategies in scalable systems like LoRA and MOE, where task-specific routing is essential (already widely used in the ML system of industries). Similarly, in speculative decoding, system-level routing frameworks like MetaSD have significant potential to enhance efficiency and flexibility.

We hope this analysis provides clarity on the representativeness of our experiments and the broader implications of our framework.

Warm regards,

Authors

Dear Reviewer v7AR,

Thank you again for your detailed feedback. We would like to address the concerns you raised in more detail and clarify our assumptions and theoretical results.

Extrinsic Randomness of Model Alignments

First, we would like to emphasize that even in Greedy Decoding, randomness arises due to uncertainties in the alignment between the target model and the drafters. Specifically, let $\alpha_{i,t}$ denote the acceptance rate of the $i$-th drafter at the $t$-th speculative round during one decoding instance. Even if the target model and drafter are deterministic in their outputs, it is practically impossible to access full information of $\alpha_{i,t}$ (i.e., exact alignment or acceptance rates) in advance.

For example,

• The number of accepted tokens by a drafter depends on the specific overlap between its output and the target model's sequence at each step, which can vary across tokens.

• Thus, before observing the outcomes, the system naturally treats the alignment as a random variable.

We model the variation of number of accepted tokens (or $\alpha_{i,t}$) in each iteration as a random sequence i.i.d. from stationary distribution with mean $\alpha_i$ in each decoding instance (Please note that this does not imply that we assume i.i.d over several decoding instances). Therefore, our assumption is built from extrinsic randomness arising from modeling the randomness of the alignment between the target model and a drafter, which is orthogonal to considering the randomness coming from the stochasticity of the model outputs which we might call intrinsic randomness.

Assumption 1 and Assumption 2 is about the extrinsic randomness which models uncertainties of how each drafter is aligned with the the target model and i.i.d assumption on the acceptance rates can be made within this scenario. And as a result, Assumption 1 can hold in Greedy Decoding scenario even if there is no intrinsic randomness of the models.

The Setting of Theorem 2

Regarding Theorem 2, we would like to clarify once more that it does not assume fixed B of LLM output. Instead, Theorem 2 analyzes the regret bound defined in each instance of decoding which is independent of the output of each generation (as explained in the follow-up response of W3).

For example, suppose we conduct 2 decoding procedures (same rule as your new example) and generate 4 tokens. Then, the performance metric should measure how well the policy followed picking the best drafter in each decoding instance which does not related to the correlation between the model output and the target sequence length. Since our regret is analyzed on each decoding instance, first instance (B=2) should be measured by how well the algorithm performs compared to the optimal policy which is simply choosing the best drafter consecutively, and the second instance (B=2) should be measured by the same criterion. With the given $B=2$, the number of rounds $\tau(B)$ is still stochastic under our modeling of the 'extrinsic randomness'. Our stopping regret objective is thus to reduce $\tau(B)$ for each instance ($B=2$ for the first one and $B=2$ for the next one), which is independent of the results of target model. Even if the outputs of the two instances are different (or the same), this does not break the Theorem under the Assumption 1 as we stated in the Appendix G.6.

Independence of decoding instances and Theorem 6

We would like to once more clarify that we do not make assumptions about the independence of each decoding output. Instead, we make assumptions on the independence of the acceptance rates for a given drafter along time for each instance of decoding. As stated in the above response (follow up response of W4), Theorem 6 is built with our additional assumption about assuming $\alpha_i$ being independent of $B$ and its conditional expectation over a given $B$ being the same for every $B$. With Assumption 2, the acceptance rates follows i.i.d. conditioned on each B and therefore, Theorem 6 remains valid under the Assumption 2.

We will revise our manuscript accordingly to further clarify our assumptions is built on the model uncertainties itself (extrinsic) rather than the randomness of model outputs by temperature sampling (intrinsic). Also, we will further clarify Theorem with more explanation on the role of assumption.

We hope these clarifications address your concerns. Thank you again for your active and constructive feedback and we are open to further discuss if there are any remaining concerns.

Warm regards,

Authors.

• Theoretical guarantee: Provides a solid theoretical foundation for multi-armed bandit-based drafter selection in speculative decoding.

• First of its kind: The first work to address speculative decoding with multiple heterogeneous drafters, tackling real-world challenges of diverse task settings.

• Extensible Design: Easily integrates with other speculative decoding frameworks, offering compatibility and flexibility for future research.

• Robustness: Demonstrates consistent performance across diverse scenarios, including in-domain, out-of-domain, and perturbed prompts, outperforming prior work in single-batch settings.

• Throughput limitation: Acknowledges limitations in heterogeneous large-batch scenarios but remains highly effective for small batch sizes.

• Novel analysis: Introduces a new form of divergence metric analysis, shifting focus from block efficiency to token-level adaptation and drafter selection.

This work lays the groundwork for further exploration in multi-drafter speculative decoding while addressing practical and theoretical challenges.