Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment

Author response (1)

We thank the reviewer for the very thoughtful and extensive review. We will address the weaknesses and questions in the following: 1. **Verification is changed:** We completely agree with the reviewer, and we hope that this is evident enough from our main text; speedup comparisons are not enough anymore as we (deliberately) give up on the guarantee, and quality checks need to be performed. We perform various checks by evaluating our resulting model on several standard tasks for LLMs, demonstrating that our model does maintain quality. We further agree with the reviewer that some alignment between the judge dataset and the tasks of interest is needed; if we remove coding examples, we observe a significant drop in performance as shown in the main text. On the other hand, we also observe that tasks such as “MMLU” and “MT-Bench” did not need further hand-curated prompts to maintain quality, as they seem to share enough overlap with the dataset already. We thus suspect that a large enough dataset should cover a large fraction of the prompt space, but we definitely agree that very “novel” tasks most likely require annotating new samples. We also want to point out however that the number of such examples is relatively small in our experience, ranging only from 10-20 examples. 2. **Open-ended tasks:** We thank the reviewer for this great insight, we did not highlight this limitation enough in the main text. Our approach relies on the fact that there are correct and wrong tokens, and it is not immediately obvious how to extend this to more ambiguous tasks such as writing, summarization, and quality of response in general. One could think about ranking responses, similar to the mechanism used in RLHF to distinguish such finer differences in quality, which would make for exciting future work. We will definitely add these points to the paper! On the other hand, we do want to highlight that MT-Bench does measure the quality of responses through a GPT-4 judge, which does include some more ambiguous tasks such as writing and summarizing, and we do not observe a strong regression there. We view our work as a first step towards more flexible verification, and we will definitely highlight this limitation in the main text. 3. **High-quality draft:** We agree with the reviewer, that a high-quality draft model such as Llama-8b that is capable of drafting longer, sensible answers is needed for our approach. Otherwise, accepting more tokens will prove detrimental. We view this limitation as less worrying for the following reasons: (1) Small models have been improving rapidly over the last year, and we expect this trend to only get stronger. This means that fast draft models of high quality will become more and more available. (2) On the other hand, flagship models will most likely continue to grow in size, making thus “better (larger) models small enough in comparison”. E.g. one could imagine Llama-70b becoming a fast enough draft model if the target were as big as Llama-”3005b”. We are hence confident that this limitation is not and will not become a crucial one in the future. 4. **Out-of-distribution:** We kindly ask the reviewer to clarify what is meant by “full results not being reported”. We have added a table in Appendix B.5 that evaluates the same judge without coding examples on the other tasks and as expected, we observe roughly the same scores (except for MT-bench, which also includes a small set of coding tasks). If that was not what the reviewer meant, let us know! We are reporting the average speedup over GSM8k, MT-Bench and HumanEval in Table 1. We have added a table in Appendix B.4 further detailing the individual speedups. As has been observed in prior work, HumanEval (a coding task) enjoys the most speedups, followed by GSM8k and then MT-Bench. 5. **Amount of prerequisites:** Again, we don’t disagree with the reviewer, very new tasks will require the annotation of some examples (as we show in the paper), but we also found that number to be very small (10-20 prompts). Other tasks such as MMLU on the other hand were even already covered well by simply having other multiple choice prompts in the training set, so the concept of “correct” and “wrong” does generalize between different tasks. We thus believe that the amount of additional data needed for new tasks will shrink significantly as the training data for the judge grows. We further believe that the very significant speedups of our technique do justify some of the overhead, especially when the domain of application is well-known in advance. Finally, we would like to stress again that we don’t want to claim that our approach represents a silver bullet. To the best of our knowledge, it is the first work in this novel direction, and we are convinced that future work can build upon it and further refine it. We are happy to further stress this in the paper, but we believe that we were already very honest about the limitations, as the reviewer also points out.

Author response

We thank the reviewer for the very positive feedback and the helpful review. We will address the questions and weaknesses in the following: 1. **Less binary tasks:** This is a great insight and we agree that we did not highlight this limitation enough in the main text. Indeed, our approach currently relies on the fact that there are correct and wrong tokens, and it is not immediately obvious how to extend this to more ambiguous tasks such as writing, summarization, and quality of response in general. One could think about ranking responses, similar to the mechanism used in RLHF to distinguish such finer differences in quality, which would make for exciting future work. On the other hand, we do want to highlight that MT-Bench does measure the quality of responses through a GPT-4 judge, which does include some more ambiguous tasks such as writing and summarizing, and we do not observe a strong regression there. In general, however, we definitely agree with the reviewer that stronger regressions are likely to be observed in such settings. We view our work as a first step towards more flexible verification, and we will highlight this limitation in the main text. 2. **Related works.** We thank the reviewer for sharing this relevant work, we have included it in Section 4.1.

Author response

We thank the reviewer for the very positive feedback. We will address the very valid point of missing related works below: **Related works:** We agree with the reviewer and apologize for this oversight on our part. We added the following works on speculative decoding to the related works section: 1. *Blockwise Parallel Decoding for Deep Autoregressive Models*, Stern et al., 2021 2. *Instantaneous grammatical error correction with shallow aggressive decoding*, Sun et al., 2021 3. *Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation*, Xia et al., 2022 4. *Accelerating Transformer Inference for Translation via Parallel Decoding*, Santilli et al., 2023 We’re happy to further expand this section in case we are missing more relevant works.

Author response

We thank the reviewer for the very positive feedback and the very helpful review. We will address the questions and weaknesses in the following: 1. **Some of the results are obvious:** We agree that it is not surprising that standard speculative decoding fails when very different models or human text are used as drafts, precisely due to the reliance on alignment. We think the more surprising or interesting part is that our method, judge decoding, does manage to accept a high number of tokens in these scenarios, as it relies now on the contextual quality. These experiments thus highlight the switch in the evaluation criterion: SD relies heavily on alignment, while JD evaluates quality. We hope this explains the inclusion of these more “obvious” experiments, as they are needed to enable later comparisons to our approach. 2. **Disjunctive criterion:** This is a great question and the choice seems to be a bit unnatural at first. We rely on the disjunctive criterion for the following reason: *If we replace a JD-rejected token with the most likely token according to the target model, we naturally end up with the disjunctive criterion.* We have also added a more detailed explanation in Appendix B.2 to clarify this further. If JD rejects a token, we replace this rejected token with the corresponding token produced by the target (just like in standard SD). Sometimes this target token is identical to the JD-rejected token. This can sometimes happen because JD is calibrated to rather err to reject tokens than to accept wrong ones. We thus “correct” the JD-rejected token with the same target token, and continue generating. Standard SD on the other hand would accept in this scenario (the proposed token is identical to the target token) and the course of sampling wouldn’t change; we would simply save computations by accepting the token already in this step, instead of adding it and re-generating draft tokens. Thus, as soon as we decide to replace rejected tokens with corresponding target tokens, we naturally end up with the disjunctive criterion. In order to make progress once a rejection is reached, we don’t see other alternatives beyond simply taking the corresponding target token. We hope this makes sense! In our experiments, however, we don’t observe this very often, JD accepted tokens are almost always a superset of the SD accepted ones. We still implement inference this way however to be most optimal. 3. **OOD performance:** The dataset largely consists of general input prompts from the Alpaca dataset [1] that were hand-filtered by us, which were not targeting specific evaluations but served to teach the judge the concepts of “correct” and “wrong” from many angles. A smaller percentage of prompts were created to more closely mimic evals, e.g. we added math questions to match performance on GSM8K, and we added coding examples to improve on HumanEval. The number of such additional prompts needed turned out to be very small, usually only 10-20 examples were needed to restore accuracy. There is also a high transfer between tasks; for ARC we included 20 prompts from its training set to teach the judge to evaluate multiple-choice answers. For MMLU (another multiple-choice dataset), there was no need to include further prompts as soon as ARC training prompts were added, as the notion of correct or wrong for MC questions of one task seems to be enough to transfer to others. In summary, there definitely is some alignment between the tasks we consider, and the dataset constructed, but we find that only a few prompts are needed for some tasks, while other tasks already end up being covered (e.g. MMLU and MT-Bench). In general, we expect the amount of new data needed to reduce significantly in the future as the dataset starts to grow and cover more and more aspects, i.e. the amount of completely new tasks (such as coding or math) will reduce quickly. 4. **Fig 2/Fig 3:** Since we cannot “speculate” with closed-source models or human text, we decided to produce full answers to all prompts and simply evaluate how many tokens are accepted until we reach the first rejection within the entire response. There was no natural choice for “number of draft tokens” and we thus decided to use full answers instead. 5. **Draft/Target switched:** This is an interesting observation! The rates are similar (6.3 vs 6.6) but slightly in favor of using Llama-405b as the draft and Llama-8b as the target, as the reviewer correctly points out. We think that Llama-405b might be better at mimicking the style of the outputs from Llama-8b the more it gets conditioned on its outputs (since it's a more capable model), compared to Llama-8b matching the style of Llama-405b. It thus seems reasonable to expect a better rate for Llama-405b as the draft model. We would be happy to discuss this further with the reviewer. We will update the paper to add this explanation. [1] *Stanford Alpaca: An Instruction-following LLaMA model*, Taori et al., 2023

Clarifications

Hi Weilin, Thank you for your interest and careful reading of our work! We want to clarify a few things: 1. **Medusa speedup:** The 6x for Medusa in HuggingFace is an upper bound, since 6 tokens could be maximally accepted if things would work perfectly with zero overhead. We don’t report a concrete number because [1] did not opensource their model and only ran it in (yet) another optimised framework and on different hardware, resulting in the 1.9x number we report. We mark it with an asterisk for these reasons. We thus don’t recommend comparing things directly using the 6x number. 2. **Reversing speedups:** We agree that at first it seems counter-intuitive that the trend reverses for the 70B target case, where Eagle speedups become worse in *gpt-fast*, while judge-decoding becomes faster. The driving factor in these speedup numbers is the discrepancy between draft and target speed, which changes in surprising ways when switching between frameworks. **In short:** 405B/8B latency gap reduces when compiling, while the gap of 70B/8B increases when compiling, thus leading to the reversed speedups. We will explain in more detail below: - **8B/405B**: 405B is extremely slow in HF, leading thus to a large discrepancy to 8B even without compiling (in contrast to the 8B/70B setup). As a consequence our approach enjoys large speedups since we have a large discrepancy between target and draft. When compiling, both 8B and 405B become significantly faster in isolation (as expected). However, 405B enjoys even more speedup than 8B (we’re not entirely sure why but HF is probably not optimised for this size) , which thus leads to a **reduction** in terms of the speedup ratio, since the discrepancy between 8B and 405B is smaller now. Medusa or Eagle heads on the other hand don’t become much faster due to compilation since they’re very small models to begin with and not much overhead needs to be removed. In their cases, only the target model becomes significantly faster, reducing thus speedup numbers. This is consistent with results reported in Eagle [2], where speedup numbers reduced from 2.8x to 1.5x when compiling and quantising. [1] also finds a speedup of 1.9x in case of Medusa when compiling, which most likely is smaller than the one in HF. - **8B/70B**: In HF, counter-intuitively, the latency gap between 70B and 8B is small. This is because both models are bottlenecked by CPU instructions at very small batch sizes. This is why SD and JD enjoy little speedups here; the draft is simply too slow compared to the target, as observed in many previous works for standard SD as well. This issue does not apply for Eagle; the very small module on top does run faster than the target even in HF and it thus dominates in this regime. When compiling the models, CPU-bound issues are removed and 8B runs significantly faster than 70B, leading to a **larger** latency gap. This leads to an improvement in the speedup numbers both for our JD approach and standard SD as well. This is consistent with results reported in the repo *gpt-fast* (https://github.com/pytorch-labs/gpt-fast), where 70B can be sped up by 2x using 8B as target, in contrast to the usual results that report 1.4x when using HF in this scenario (see e.g. [5]). Recall that in case of 405B, the target improved significantly more due to compilation than 8B did, which is why speedup numbers went down in this case. For Eagle on the other hand, again, compilation mainly benefits the target since the draft was already fast, leading thus to a reduction in speedup. Similarly, [1] find a speedup of 1.45x for Medusa for 70B, which again most likely is significantly smaller than speedups in HF (e.g. [3] reports a speedup of 2.83x in case of Vicuna-13B, there is unfortunately no larger model reported). 3. **Medusa for 70b**: This is a good point but we’re not aware of a publicly available version for Medusa in case of the 70B version. The same work by NVIDIA [1] that we cite for the 405B version does show speedup numbers also for 70B, which roughly are 1.45x when compiled. As expected based on prior work, the speedup is worse compared to our results for Eagle. We’re happy to further include this number too if this is helpful. Thank you for pointing out the useful baseline from [4]. We hope to include it in the next version of our paper. We’re happy to rectify any mistake if there are better known speedup numbers when compiling! [1] https://developer.nvidia.com/blog/low-latency-inference-chapter-1-up-to-1-9x-higher-llama-3-1-performance-with-medusa-on-nvidia-hgx-h200-with-nvlink-switch/ [2] *EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty*, Li et al., 2024 [3] *MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads*, Can et al., 2024 [4] *Fast inference from transformers via speculative decoding*, Leviathan et al., 2023 [5] EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees, Li et al., 2024