SpecTr: Fast Speculative Decoding via Optimal Transport

Summary

The paper proposes SpecTr, a new speculative decoding framework for efficient Transformer and LLM decoding. SpecTr uses a combination of a small, fast model for generating draft samples and a larger, more accurate model for scoring and validating them in parallel. Compared to the prior works, the paper presents a formulation of the speculative decoding process through the lense of optimal transport theory and demonstrates the potential for further acceleration by generating multiple drafts in parallel along the batch axis. The authors offer an approximate solution to efficiently solve the optimal transport problem at scale, resulting in significant latency improvements compared to baseline methods and previous speculative decoding approaches.

Strengths

The paper presents theoretical justifications for the methods that it proposes. Furthermore, it is the first attempt to incorporate parallelization along batch into the speculative sampling framework for better efficiency.

Weaknesses

1. While the paper proposes theoretical justifications for the proposed method, it rather lacks proper evaluation. Since accelerating decoding processes is a very practical area, the author should provide a more thorough analysis of the end-to-end latency and text generation performance. In particular: * (a) The paper lacks an evaluation of the impact of the proposed method on text generation performance. Tables 2 and 4 in the paper only focus on latency, and there is no comparison in terms of text generation quality (e.g. measured by metrics like BLEU scores, etc.) compared to the baseline or prior speculative methods. Latency numbers without these performance metrics make it difficult to assess the effectiveness of the proposed method. * (b) The number of decoded tokens per serial call (in Table 2) is not directly indicative of the actual latency and does not necessarily represent the actual speedup. That is, it serves as a proxy measurement that may or may not yield the same degree of latency improvement. This is because there are factors such as framework overheads [1], hardware utilization, and others that can potentially impact the latency. For instance, it is difficult to state if processing 3 tokens (as in Table 2) in a single call is truly 3 times more latency efficient than processing them sequentially in 3 separate calls. Therefore, the claimed 3X speedup, and the additional 1.36X improvement over speculative decoding mentioned in the abstract, can be misleading. * (c) It is unclear whether the proposed method introduces any additional run-time overhead compared to the prior speculative sampling method. If there is indeed extra overhead involved, it could potentially reduce the gap between the two methods, making the improvements presented in Table 2 less significant (thereby making the 1.36X improvement over the prior method misleading as well). * (d) The latency overhead of running smaller models concurrently remains uncertain, since Table 2 only presents the efficiency of running the large model. For instance, a longer sequence length (L) can decrease runtime costs for the large model by increasing the number of tokens decoded per call as stated in the table. However, at the same time, it also increases the small model's cost as more tokens are generated but rejected/discarded. The paper lacks discussions around this point. While the authors presented the latency overhead of 10-15%, it is not clear under which particular setting it was measured and how the hyperparameters (K and L) affects this value. * (e) Taken together, the authors should provide the end-to-end inference latency and the text-generation quality/performance measure to make the evaluation convincing. 2. Including a concluding paragraph in the paper would enhance its professionalism. The authors should at least wrap up the paper with a one-paragraph summary. [1] The Framework Tax: Disparities Between Inference Efficiency in Research and Deployment, https://arxiv.org/pdf/2302.06117.pdf

Questions

N/A

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Please see the weakness section

Summary

The paper studies speculative decoding, a technique to increase inference efficiency in a large autoregressive model by sampling a set of candidate tokens (a *draft*) from a smaller (thus, faster) model, which is then scored according to the conditional distribution of the original larger model. The scores of the larger model are obtained in parallel, thus potentially resulting in a significant speedup depending on how well the smaller model's conditional distribution is close to the larger model's. The main contributions are twofold. First, the authors cast the speculative decoding problem as a maximal coupling problem in optimal transport. Then, they propose to use multiple drafts (i.e. multiple sets of tokens completing the context), which is solvable in exponential time in the number of drafts. Thus, a novel algorithm is proposed that runs linearly in the number of drafts and provably solves the problem up to a $(1-1/e)$ factor of the optimal acceptance probability.

Strengths

1. To my knowledge, translating the speculative decoding problem in the context of optimal transport is not very surprising but novel. Furthermore, by connecting the problem to the field of optimal transport, the paper opens new possibilities for research, due to the maturity of the theory of optimal transport, which the authors actually use to derive an extension of the original speculative decoding algorithm to allow for multiple drafts and discuss its feasibility in terms of algorithmic complexity. 2. The proposed algorithm to approximately compute the transport plan (Algorithm 2) nicely follows the problem casting performed earlier and is backed by a sound theory (Theorem 2) describing how the solution (in terms of acceptance probability) compares to the intractable case of exact solution in exponential time. I have quickly skimmed through the appendix and although not an expert in OT, the theory seems correct.

Weaknesses

1. The idea of extending the speculative decoding algorithm in [15] to multiple drafts is simple and elegant, but sometimes the authors add some unnecessary details. For instance, it seems trivial and very intuitive that by having multiple drafts, the acceptance probability increases with the number of drafts (hence Lemma 1 seems unnecessary). 2. The authors test their proposed method solely on LM1B. I would have liked to see a larger set of experiments, covering at least a more significant subset of those performed in [1]. 3. Algorithm 1 is wrong, as it always returns at line 7. I guess the authors want it to return at line 5 as well. Also, some quantities are not defined yet when Alg.1 is introduced (such as $\mathcal{X}$), which complicates a bit the readability. Minor: 3. In Figures 1,2,3 the size of ticks, labels, and lines should be increased. 4. Typos: line 251: "Gvien". Line 217-218 the sentence "To control the probability of accepting an $x \in \Omega$ with probability larger than $q(x)$." does not seem to be grammatically correct. Line 44 "a several contexts". Line 120: "an discrete" [1] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. arXiv preprint arXiv:2211.17192, 2022.

Questions

None.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors address the limitations. I do not see any potential negative societal impact.

Summary

This paper presents speculative decoding, wherein a smaller language model is used to approximately sample from a large one, akin to a type of accept/reject MCMC sampler. The idea is that it is slow to sequentially sample from a large model of interest, but joint probabilities can be computed in parallel across the time dimension. At the same time, we could have access to many parallel copies of a smaller not-as-good model, and it's quicker to generate e.g. 100 tokens from the small model than to sample from the large one. We generate e.g. 100 tokens for K such smaller models, and accept one of the K sequences with some probability. The paper is mostly a theory paper that gives bounds on the transport cost, between the large model distribution and the sampling distribution resulting from the accept/reject sampling composed with proposing from the smaller models. On the plus side, the paper is a theory paper that I think has some nice implications for practitioners. I would love to see some more empirical work exploring this. The quality of the math / theory / bounds is high and the authors did a nice job breaking down the OT parts for the typical ML language modeling reader. On the minus side, while a light set of experiments is totally fine for a theory paper, I believe most readers of the paper would benefit more if there were also some initial results of the actual sample quality from the presented algorithm. There seem to be results on latency and accept/reject probability (and there's some theory connecting accept/reject probability to closeness of sampled distribution to desired large model distribution), but there seem to be no direct results evaluating the samples resulting from the new algorithm. In short, I tentatively accept but think the minus needs to be addressed and will change the score if there is no discussion or clarification on this. Note that light evaluation is totally fine since the theory is good. And it is even okay if the quality of the samples according to any of the metrics is possibly below competitive, since there is always room to improve these algorithms once they are established.

Strengths

- novel algorithm for quicker sampling making use of smaller models and MCMC-like concepts to sample from a larger model - nice framing in terms of optimal transport that I think the typical ML + LM reader could take something positive away from - rigorous characterization theoretically of the algorithm in terms of transport cost bounds, accept/reject rates, etc.

Weaknesses

- Table mismatch: e.g. referenced table 8 in main text but no table 8 in main text or appendix - Most noticeably, there seem to be no direct reports of quality of samples from the resulting approximate algorithm. There are tables for latency and for accept/reject rates (the latter is in appendix but could probably go in main text). And there is theory characterizing that the transport cost will not be too bad, etc. But, there are no direct metrics quantifying the original large model samples vs the proposed algorithm's samples (those produced by using the smaller models' sampling and larger models' scoring). Seems like any of the basic LM metrics would be warranted here Glad to be corrected if I have misunderstood. Minor: - Small type "Gvien" in source line 251 of the PDF.

Questions

See Strengths/Weaknesses

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper proposes a novel and efficient decoding algorithm for autoregressive large language model, SpecTr, which is an extension of speculative decoding. Given a large model M_b, it can only output one word at a time when it decodes, however, the cost of a serial call is quite expensive. The previous method speculative decoding can alleviate the slow decoding problem, which uses a much smaller model M_s to decode a single segment of length L and allows M_b to calculate its likelihood in parallel. The optimal transport plan between the distributions of the two models is used to determine whether to accept a certain part of the segment, which makes sure the samples has the same distribution as M_b. This approach has the advantage of reducing the number of serial calls of M_b, improving decoding efficiency. However, if the acceptance probability is too low, the reduction may not be significant. The SpecTr algorithm allows to utilize K i.i.d. segments sampled by M_s to enhance acceptance probability. Concretely, the previous optimal transport problem is modified to transport the K-product distribution of M_s to the distribution of M_b. Since solving the new problem has an exponential cost with respect to K, the paper introduces the k-sequential selection algorithm with a 1-1/e approximation ratio with an appropriate cost. In numerical experiments, they use an M_b model with 97M parameters and an M_s model with 6M parameters tested on LM1B, speculative decoding achieves an average of decoding 2.3 tokens per serial run of M_b, while the proposed method in this paper can decode 3 tokens, showing a significant improvement.

Strengths

The paper is well-written and self-contained, making it accessible even to someone like me who is not familiar with the related work on language model decoding. The proposed algorithm, SpecTr, combines the practical need for parallelization of large models in real-world scenarios. Moreover, it is mathematically concise, intuitive, and insightful. The paper also provides a basic theoretical analysis of the properties of the algorithm. The numerical experiments show a significance improving of their methods. I think this is a good submission.

Weaknesses

The numerical experiments on real datasets seem insufficient, as they only compare with speculative decoding when K=1, which may not be fair: using larger K would require more computational resources. It would be worth considering if there is a more equitable way to compare the methods, e.g. extend speculative decoding to use more resources in some crude way.

Questions

The first line of Algorithm 1: Input: ... "X~iid p" is confusing, since you have only one sample here, "iid" is unnecessary.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

No