Multi-Draft Speculative Sampling: Canonical Decomposition and Theoretical Limits

Authors: Ashish J Khisti, MohammadReza Ebrahimi, Hassan Dbouk, Arash Behboodi, Roland Memisevic, Christos Louizos
Primary area: foundation or frontier models, including LLMs
Keywords: speculative decoding, multi draft speculative sampling, large language models, weighted importance sampling, optimal transport

Abstract:
We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models.  At each step, a  token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accepting one of the input tokens) can be cast as a solution to a linear program. In this work we show that the optimal scheme can be decomposed into a two-step solution: in the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token.  For the case of two identical draft models we further 1) establish a necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability.  Our theoretical analysis also motives a new class of token-level selection schemes based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number of scenarios.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1729566472025

thread_type:
official_review

<thread_object>
creator: reviewer_yfzp

<content>
summary:
This paper studies the problem of multi-draft speculative sampling, where the optimal scheme can be cast as a solution to a linear program. The key idea behind this work is that this optimal scheme can be decomposed into a two-stage solution. Based on this idea, the authors provide some theoretical findings towards the two-stage solution. Motivated by the analysis, the authors propose a new class of token-level selection scheme based on weighted importance sampling, which shows consistent improvement in various scenarios.

strengths:
1. The introduction of the two-step scheme are simple and intuitive, making it easy to understand.
2. The authors provide a solid theoretical analysis for token-level optimal draft selection.
3. The truncated LP algorithm is novel to the field of speculative decoding and theoretically better than SpecTr.
4. Extensive experiments, including different datasets, different decoding temperature and different draft models demonstrate the effective of proposed method.

weaknesses:
1. I think the contributions of this work is a little weak. Given that speculative decoding shows poor performance with batch size larger than 1, the multi-draft sampling is usually not an appropriate option for speculative decoding. I think the authors should emphasize the contributions to real-world speculative decoding applications. 
2. While I appreciate that the authors add some examples for better explanation, I think the authors should add a clear notation section with distinct description.
3. For the truncated LP draft selection algorithm, there exists an operation of sorting, which takes a complexity of $O(n\log n)$ and may influence the final speedup. The authors should report its real-world time consuming and compare to vanilla multi-draft selection, especially the empirical results on models with large vocab size, e.g., Llama 3.1 with vocab size 128k. Besides, for multi-draft scenarios with $K > 2$, the authors propose a multi-stage manner, which may affect the final speedup as well.
4.  I think the authors should conduct experiments on more benchmarks (e.g. HumanEval [1], GSM8K [2] and MT-bench [3]) for a fair comparison. Besides, the authors should report the real-world speedup to demonstrate the effectiveness of proposed method.

I know that conducting a wide range of evaluation experiments is costly and time consuming, and I will increase my score if the authors clearly address my concerns.

[1] Chen, Mark, et al. "Evaluating large language models trained on code." *arXiv preprint arXiv:2107.03374* (2021).

[2] Cobbe, Karl, et al. "Training verifiers to solve math word problems, 2021." *URL https://arxiv. org/abs/2110.14168* (2021).

[3] Zheng, Lianmin, et al. "Judging llm-as-a-judge with mt-bench and chatbot arena." *Advances in Neural Information Processing Systems* 36 (2023): 46595-46623.

questions:
1. Can proposed truncated LP combined with SOTA speculative decoding (e.g. Medusa [4] and Eagle [5]) methods for better performance?
2. Why the authors do not discuss and compare a multi-draft selection method RRS [6]?
3. Could you please discuss the limitations of this paper?

[4] Cai, Tianle, et al. "Medusa: Simple llm inference acceleration framework with multiple decoding heads." *arXiv preprint arXiv:2401.10774* (2024).

[5] Li, Yuhui, et al. "Eagle: Speculative sampling requires rethinking feature uncertainty." *arXiv preprint arXiv:2401.15077* (2024).

[6] Jeon, Wonseok, et al. "Recursive speculative decoding: Accelerating llm inference via sampling without replacement." arXiv preprint arXiv:2402.14160 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698723090

thread_type:
official_review

<thread_object>
creator: reviewer_v7r6

<content>
summary:
This paper studies the optimal token-level draft selection for a speculative sampling. In particular, speculative sampling is a method for accelerating autoregressive sampling by using smaller "draft" models to generate candidate tokens and a larger model to determine which tokens to accept. The key problem studied in this paper is how to propose tokens to optimize the probability they will be accepted. Importantly, increasing the acceptance probability leads to less calls to the larger model which in theory could increase the tokens per seconds. The key theorem of this paper is to characterize the necessary and sufficient conditions when two tokens are generated by the draft model. Interestingly, unlike the one token case, this paper shows that one can reach probability 1 acceptance even if the draft and large models have different token distributions. Finally, the paper ends with (i) heuristics to practically leverage the proposed sampling scheme, e.g., truncating the token distributions (ii) a heuristic scheme for leveraging the proposed importance sampling to larger sequences and (iii) empirical evidence of the effectiveness of the approach for K=2 and larger.

strengths:
This paper provides a strong theoretical contribution to a very important topic. Speculative decoding serves as a refreshing and novel direction for accelerating LLM inference. This work puts the multi-draft model selection problem on better theoretical footing.

The results seem to be sound from my reading and no obvious errors were found.

weaknesses:
By and large the biggest complaint I have with this draft is the background exposition. Prior to reading this I was not familiar with speculative decoding. While I understand space is tight, many key concepts such as the role and form of accepting probability in the speculative decoder were not clearly explained. This is a *key* aspect of the work, and the draft would strongly benefit for giving it more treatment.

To be honest, I couldn't understand anything the first time I read the paper and had to go back to the original speculative sampling paper [1] before I could piece things together.

Nevertheless, after reading, I now understand that this paper is targeting the acceptance criteria as it determines the number of candidate tokens used via a Bernoulli like-process. 

This plus the strong theoretical analysis are enough for me to lean towards acceptance.

[1]Accelerating Large Language Model Decoding with Speculative Sampling,  https://arxiv.org/abs/2302.01318

questions:
- Could you give some intuition about the conjecture in remark 2. It says that the exponent over the sum of the p-weights should match K (the number of tokens to generate per round). It's not at all clear to me why these should be so tightly coupled and why other lower order terms wouldn't appear.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730766925064

thread_type:
official_review

<thread_object>
creator: reviewer_6esd

<content>
summary:
The paper studies the problem of multi-draft speculative sampling and makes new observations and provides a new algorithm. Of the new observations, Theorems 2 and 3 standout which characterize the optimal acceptance rate for two drafts and further show when the acceptance rate is one. The new algorithm exploits certain structure in the optimal transport problem and provides a clever approximate solution.

strengths:
The paper is well written and easy to follow. Theorems 2 and 3 are novel and the experimental results are convincing.

weaknesses:
See questions section below.

questions:
- In Figure 3, experiments are conducted for draft temperatures > 1.0. However typically, target model temperature is often set to 1.0 or less than 1.0. Given this, I was expecting draft temperatures also to be <= 1.0. I was wondering how these draft temperatures were chosen?

- Would it be possible to say how the gap between the proposed approach and previous approaches vary as the target model temperature varies from 0.0 to 1.0?

- Theorems 2  and 3 provide a nice characterization of acceptance probability. However, given two distributions p, q, evaluating the condition seems to require 2^{vocabulary size} number of computations. Is there a faster way to evaluate these quantities?

- There are some approximate ways to solve the optimal transport by Sinkhorn methods e.g., https://amsword.medium.com/a-simple-introduction-on-sinkhorn-distances-d01a4ef4f085. I was wondering if these ideas are applicable here?

- Architectures in ML, typically refer to neural architectures. Given several works which propose different drafting methods for speculative decoding, I was thinking "canonical architectures" refers to some new drafters. It might be good to clarify this distinction early on in the paper.

- Can you please expand on what token rate means? Is it the number of tokens / second?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730875838621

thread_type:
official_review

<thread_object>
creator: reviewer_avvf

<content>
summary:
Previous papers showed the optimality for a token selection rule when one token was chosen from the draft model. The authors do the same for the case when two tokens are selected. The authors decompose the problem of token selection into two steps --- first select one of K tokens using an importance weighted sampling technique, and then using the optimal single token selection proposed in previous papers. Finding the optimal token selection rule then reduces to finding the best importance weighted sampling technique.

The authors then show a necessary and sufficient condition for the existence of an optimal token selection rule where two tokens are selected from the draft. They then present an algorithm to solve the resulting linear program faster than the O(n^2) standard solution.

The authors compare the Importance sampling algorithm with SpecInfer on various tasks, and show that IS comes closer to optimal than SpecInfer.

strengths:
The results in the paper are quite original and significant, to my knowledge, which is a bit limited. The experiments show that the proposed algorithm works well in practice.

weaknesses:
I would have liked the presentation in Section 4 to be a little clearer - request you to write out the full linear program, and then explain the truncation a bit more - what is the truncated program computing? 

In the experiment section, please include results for Temperature < 1, which is much more common in practice.

questions:
I would like to see a comparison to the results in https://openreview.net/forum?id=9KxnxWOBA5
</content>
</thread_object>

</thread>
<thread>
creation date:
1732294543674

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for their comments. We have uploaded a revised version of the paper, and we summarize the main changes here:

- **New Experimental Results**: We have added a new experimental result (Figure 4 in Section 5 of the main paper) where the temperature of the target model is varied between 0.2 and 1.0, while the temperature of the draft models is set to 1.0. In Appendix J1, we present an additional experimental result where both drafts have the same temperatures, and in Appendix J2, we present an experimental result involving three draft models. These results complement the experimental results in Figure 3 in the main paper.

- **Main Result**: Our result in Theorem 1 is not restricted to a product distribution over the tokens; it applies to an arbitrary distribution. We have updated Remark 1 in the main paper and Appendix B.1 to indicate this. In particular, the decomposition result also applies if the draft tokens are sampled without replacement.

- **Truncated Linear Program**: We have revised Section 4 to better explain the truncated linear program. We also present the computational complexity of the truncated LP in Appendix G, as well as another variant that avoids the sorting operation in the truncated LP. 

- **Review of Speculative Sampling**: We provide a brief review of speculative sampling in Appendix A. 

- **Notations**: We have added a section in the Appendix (Section K) summarizing the mathematical notations used in the paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers the problem of multi-draft speculative decoding and derives a decomposition of the optimal solution as an importance sampling step followed by a token-level verification step. This decomposition helps the authors focus on optimizing a valid importance sampling distribution in the context of speculative decoding. For the case of two drafts, Theorems 2 and 3 fully characterize the optimal acceptance rate and further show when the acceptance rate is one. In particular, it is shown that the acceptance rate could be one even if the two distributions are not the same, which is unlike the guarantees in prior work. The decomposition also leads to a new algorithm that exploits the structure of the optimal transport problem and provides a clever approximate solution. Overall, the paper makes significant theoretical and experimental progress in multi-draft speculative decoding that deserves to be highlighted in the conference. Congratulations to the authors! 

P.S. It would be nice to discuss and relate to the concurrent work on deriving the optimal acceptance rate of multi-draft speculative decoding (https://openreview.net/forum?id=9KxnxWOBA5) in the camera-ready version of the paper.