Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment
Authors: Gregor Bachmann
Sotiris Anagnostidis
Albert Pumarola
Markos Georgopoulos
Artsiom Sanakoyeu
Yuming Du
Edgar Schönfeld
Ali Thabet
Jonas K Kohler
Venue: ICLR 2025 Oral (2025)
Keywords: LLM inference
speculative decoding
Primary area: foundation or frontier models, including LLMs

Abstract: The performance of large language models (LLMs) is closely linked to their underlying size, leading to ever-growing networks and hence slower inference. Speculative decoding has been proposed as a technique to accelerate autoregressive generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based on their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as GPT-4o, as well as human text cannot achieve high acceptance rates under the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelmingly likely when solely relying on alignment of draft and target.
We thus ask the following question: Can we adapt verification to recognize correct, but non-aligned replies? To this end, we draw inspiration from the LLM-as-a-judge framework, which demonstrated that LLMs are able to rate answers in a versatile way. We carefully design a dataset coined TokenCourt to elicit the same capability in the target model by training a compact module on top of the embeddings to produce ``judgements" of the current continuation. We showcase our strategy on the Llama-3.1 family, where our 8B/405B-Judge achieves a speedup of $9\times$ over Llama-405B, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up to $141$ tokens/s for 8B/70B-Judge and $129$ tokens/s for 8B/405B on $2$ and $8$ H100s respectively.

Published as a conference paper at ICLR 2025
JUDGE DECODING: FASTER SPECULATIVE SAMPLING REQUIRES GOING BEYOND MODEL ALIGNMENT
Gregor Bachmann 1.8
Sotiris Anagnostidis
Albert Pumarola'
Markos Georgopoulos
Artsiom Sanakoyeu'
Yuming Du
Edgar Schönfeld
Ali Thabet
Jonas Kohler
ABSTRACT
The performance of large language models (LLMs) is closely linked their derlying size, leading to ever- ever-growing networks and hence slower inference. Spec- ulative decoding has been proposed generation, leveraging a fast draft model to propose candidate tokens, which are then verified in parallel based their likelihood under the target model. While this approach guarantees to reproduce the target output, it incurs a substantial penalty: many high-quality draft tokens are rejected, even when they represent objectively valid continuations. Indeed, we show that even powerful draft models such as oo, as well human text cannot achieve high acceptance rates un- der the standard verification scheme. This severely limits the speedup potential of current speculative decoding methods, as an early rejection becomes overwhelm- ingly likely when solely relying on alignment of draft and target.
technique to accelerate autoregressive
We thus ask the following question: Can we adapt verification to recognize cor- rect, but non-aligned replies? To this end, we draw inspiration from the LLM- as-a-judge framework, which demonstrated that LLMs able rate answers in versatile way. We carefully design dataset to elicit the same capability in the target model by training compact module on top of the embeddings produce "judgements" of the current continuation. We showcase our strategy on the L.1ama-3. family, where our 4b/405B-Juch 405B-Judge achieves speedup of 9x over amalmma444 40BB, while maintaining its quality on a large range of benchmarks. These benefits remain present even in optimized inference frameworks, where our method reaches up 141 tokens/s for 83/ -Bdgge BB/405B on 2 and 8 H100s respectively.
Judge and 129 tokens/s for
8bb Standard Verification:
To find Alyana's a age, we need to aubtract from
Question:
Chenny la 10 years old. Alyana younger Chenny.
-ab/
Judge Verification:
How old Anne if she is 2 years older Alyana?
To rind Alyana's a age, us need to subtract 4 fron Caenny's age. Chenny's 's age: 10 years
Alyans's age:
years
rind Anne age, years Anne' "s age: 5
add Alyana's age.
years old.
Figure I: Standard speculative decoding versus our judge decoding strategy for Llama- 3. 1 8B as draft and Llama 3.1 405B as target. Accepted (rejected) tokens highlighted in green (red).
"Wrr done during an internship Meta GenAI. Meta GenAI. 3ETH Züiich.. Correspondence gregor gregorbßethz.ct

---

Published as a conference paper at ICLR 2025
INTRODUCTION
Large language models have transformed the field of natural language processing in recent years, displaying astounding capabilities across various tasks (Radford et al. 2019 OpenAl ct al. 2024). The performance of these models is closely tied to their underlying size, with bigger models often achieving significantly better results across benchmarks (Kaplan et al. 2020: Hoffmann et al. 2022 For example, Meta recently released their largest and best model to date with Llama boasting an enormous parameter count of 405 billion (Dubey et al. 2024).
405B,
While offering great performance, such big models require a vast amount of resources to be de- ployed, and inference efficiency starts to pose critical problem. Due to the autoregressive nature of decoding coupled with the large parameter count, token generation becomes a memory-bount pro- cess, especially at small batch sizes (Shazeer 2019 [vanov et al. 2021 Pope et al. [2023). To speed up inference in such a setting, speculative decoding (SD) has been proposed (Stern et al. 2018 Xia al., 2023 Chen et al. 2023 Leviathan et al. 2023), a technique that leverages the fact that pro- cessing several tokens in parallel comes at no additional latency cost in the memory- bound regime. More concretely, small but fast draft model produces sequence of M candidate tokens, which are then verified in parallel by the model of interest, usually referred to as the targel model. In standard SD, the target model accepts a candidate token if it assigns the same or higher probability given the context, otherwise biased coin is flipped. If at least one candidate token is accepted, inference time is reduced as the large model needs to be called only once to produce multiple tokens. As shown in Chen et al. (2023), y such strategy provably preserves the distribution of the target model, while achieving significant speedups.
Relying on target probabilities guarantees the same output, but consequence, a token is solely judged based on its alignment with the target model, and not by its inherent contextual quality. As a consequence, current approaches set the number of candidate tokens small numbers M E (5, 7), as an early rejection becomes overwhelmingly likely when drafting more. On the other hand, the quality of "small"" language models (and thus the quality of candidate tokens) has been rapidly im- proving recently. CPT-4o and the recently introduced GPT-4o-ming ni show strong performance across many benchmarks, with OpenA/ actively recommending these models over the more expen- sive GPT-4. Similarly, Ll ama 3- has achieved bigger gains in performance over its prior iterations, compared to the larger Llama- -3. 1 708 (Dubey et al. 2024). highlighting as well that small models are catching up. Phi-3-ming n is another small model at "only" three billion param- eters that despite its size manages to match the performance of C-TT33 3. (Abdin et al. 2024).
While draft models are rapidly improving and providing increasingly high-quality answers, alignment-base based verification fails to reflect this progress, still rejecting tokens that are not perfectly aligned with the target response (see Fig. [ for example). Motivated by this insight, we explore the following question in this paper:
Can we adapt verification to assess token quality rather than alignment?
We draw inspiration from the LLM-as-a-judge framework, where LLMs are used to judge the qual- ity of other model responses to user questions (Zheng et al. 2023). These judgements exhibit very strong correlation with human ratings, making this a cheap and scalable approach for model quality evaluation. Interestingly, LLM-judges display the ability to rate answers in versatile way, allowing them to appreciate correct but potentially unaligned responses. To equip the target model with sim- ilar capabilities, we design a small dataset consisting of correct and wrong replies to user questions. We create a diverse set of responses from several models and precisely annotate the location of the mistaken tokens. We then leverage the powerful target embeddings to train a small module with the objective of predicting the correctness of given token, mimicking the LLM-judge mechanism.
In summary, we make the following contributions:
We demonstrate through series of experiments how the decision mechanism in speculative decoding rejects many high quality tokens, identifying a key limitation of the technique.
We adapt verification using ideas from L.de-, eliciting the same versatile rating capability in the target by adding a simple linear layer that can be trained in under 1.5 hours.
Using Llama 3B/70B-Judçe, decoding, achieving an unprecedented 129 tokens/s, while maintaining the quality of 1.1 11aam-400 405B on a range of benchmarks.
approach obtains speedups of 9x over standard
2

---

Published as a conference paper at ICLR 2025
2
RELATED WORKS
Speculative decoding has been used and extended in a range of works, leading to significant speedups across many model families and datasets. Several ideas for draft models have been
plored in the literature. Early papers rely on specialized draft models (Sun et al.) 2021 Xia et al., 2023) smaller versions of the target model (Chen al. 2023 Leviathan
al.) 2023). usually trained using the same data and learning protocol. Another line of work uses shallow networks on top of the target embeddings as a drafter with the goal predict multiple tokens into the future (Stern et al. 2018]] Cai et al. [2244] Li et al. 2024a b Zhang et al. [024aa Wertheimer et al. 2024] Ankner et al. 2024 Gloeckle et al. 2024 Bhendawade et al. 2024). Other approaches use a sub- network of the target model, e.g. [Schuster et (2022)) Zhang et al. (004b)); Elhoushi et al. (2024);
Liu et al. (2024b a skip a percentage of the layers to produce candidate tokens. Other architectures have also been explored: |Wang al. (2024) develop a SD variant for Mamba models (Gu & Dao 2024). while Christopher et al. (2024) explore diffusion-
language models as drafters. Other components of the process have been investigated as well; Huang et al. (2024); [Lu et al.] (2024c) analyze the number of draft tokens M with the goal of learning to choose it dynamically. Kim et al.) 2023) take a similar approach aiming measure the uncertainty of the draft model, allowing the target to take over when needed. |Monea et al. (2023): [Bachmann & Nagarajan (2024) on the other hand explore parallel decoding without draft by conditioning on "look-ahead" tokens or using Ja- cobi iterations (Santilli et al. 2023 Fu et al. 2024), allowing the target to produce several tokens in one step.
Many works have aimed to improve the acceptance rates in SD: Zhou et al. (2024) encourage higher alignment by finetuning with distillation loss. [i et al. (2024)) Cai et al. (2024); Ankner et al.) (2024); Miao et al. (2024); Chen et al. (2024) construct token trees out of the top- k predictions in various ways and verify them in parallel using areeaatention, covering thus larger space of token combinations. Other methods propose to exchange more information between draft and target: Du ctt al. (2024) allow the draft to access the key-value cache the target, while in S et al. (2024) the draft is further conditioned either encouraging better alignment with more information or producing more guesses in parallel. This is different from our work, which seeks to change the verification scheme itself.
target activations. All these methods improve acceptance rates by
VERIFICATION IN SPECULATIVE DECODING
3.1 BACKGROUND
Speculative decoding. Denote by LLMang and LLMatan the target and draft model respectively. We use V (1,..., , V} for the vocabulary. Let M E N represent the number of candidate tokens, m. the number of accepted tokens and 8 E V the context. Let us denote by
(t1, Pi)...., (fm, (4m,Pm) = LLM(") (s)
(1)
an autoregressive sampling of m tokens from LLM given context s, where t t,...., fm E V are the sampled tokens and p; E IR. the corresponding softmax probabilities. Further, we denote by
PI.--. Pm LLIM(I,..... (mmi
running the (parallel) forward pass of LLM on tokens f1 f,,...,tm. fm. Notice that this produces one more probability vector Pm++1 as we now also process token f,..
In SD, the draft model autoregressively produces M candidate tokens given the current context s using any sampling scheme (but usually greedy),
(c1, (cM, = (MM(M) (8)
(3)
where .I...., CM are the sampled candidate tokens and qi E R. for 1,..., M are the corre- sponding probability vectors over the vocabulary. The probability of token under the draft model thus [c(.... The target model then processes these tokens in parallel, resulting in probability
CM. vectors Pt PM+1 LLMaar (cI...., CM; 8). Rejection now works follows:
P:(c.]
Accept if all previous tokens are accepted and e; for e, U (10 1))
(4)
[
3

---

Published as a conference paper at ICLR 2025
Human-test/ Uema 4050
Standard-SD
Judge Decoding
85% rate
Human text/ Lomm 4030 Judge
Number of accepted tokens
40
12.3
Acceptod
30
Wikipedda-summarree
20
Uama-4450 Ubme 4059 Liema Bb
6.3
10
Accepped
10
20
30
40
50
Number draft tokens
GSMBK
Figure 2: Left: Average number generated tokens as a function of the number of draft tokens for Llama- 83/4 40BB with standard and judge verification. Right: Number of accepted tokens high-quality human text (top) and for both 8B/405B and 405B/ / 8B (bottom), both standard SD.
In words, a candidate is accepted if the probability under the target model is even larger. If the probability smaller, a stochastic decision is made according to the discrepancy between the probabilities. Crucially, one is always guaranteed to produce at least valid token: P1 is solely a function of the current context s and can thus be used to sample a token according to the target distribution. Similarly, if is the first rejected token, Finally, when all candidate tokens are accepted, an extra token can be sampled from pM+1- The accepted tokens are then added the current context and repeat the steps until completion.
sample correct token from p:-
Number of draft tokens. An immediate question comes to mind when examining speculative decoding: How many draft tokens should one choose for optimal speedup? On the one hand, the draft model produces good tokens, one would ideally want to draft high number of candidate tokens M to avoid invoking the expensive target model too many times. On the other hand, many candidates end up being rejected, The ideal M thus heavily depends on the acceptance rate, which in turn naturally depends on the verification scheme. In Fig. we plot the average number of accepted tokens
wants to avoid spending the unnecessary drafting time.
function of M for the model pair I1 ama 1aneaa-33.0 3. 1- 8B and Llama-3. 1 1-405B evaluated on MT-Benck (Zheng et al., 2023) and GSM8K (Cobbe al.) 2021) (yellow curve). We observe that the number of accepted tokens quickly saturates as function of M and the acceptance rates thus decrease rapidly. As result, choosing large number of draft tokens M solely calls the draft model more in vain, leading to inefficient inference overall. This is the reason why prior work is limited largely to M < 7.
3.2 LIMITATIONS OF STANDARD VERIFICATION
Rejected tokens. What types of tokens get rejected in such setup? In order to obtain an intu- ition, explore the behaviour of SD on several benchmarks including GSM8K, -TBBench and HumanEval (Chen et al. (02T). We use Llama-8 as the draft and 4llma-4055B as the target model. While there a significant discrepancy in terms performance between these two models, it is worth highlighting that the draft model achieves competitive scores all these tasks. Higher acceptance rates would thus not necessarily reduce the quality of the output on many these ex- amples. In fact, large number of draft answers could be accepted they are, especially those addressing relatively simple queries.
Notably, even in instances where the draft model produces entirely accurate solutions, the target model frequently rejects numerous tokens due to the stringent nature the verification process. This rejection occurs despite the correctness of the solution, the target model seeks alignment with its own response rather than contextual accuracy (Liu et al.) 2023). As illustrated in Fig [ a correct answer can be rejected after only tokens, underscoring the potential for relaxed verification

---

Published as a conference paper at ICLR 2025
4P Usin 485o
8-4o/ 40-b Standard so Verification:
CFT44 Lama-4050- 4od0-
Usec Queation:
Here's Python function that takes
MT-
fuaction Python takea a liat of integere, every even integer and deletea every add inseger.
40/ Judge Verification:
23
Heref' a Eychon Function that takes integers, repeats every integes, deleces evecy intager:
GSM6K
daf modify 11st
result
for num in numbers:
21
numbers
if num O:
HumanEval
result.ene eextnnd ( (nun, nun)) odd, numbers skipped result.
26
Figure 3: Left: Standard SI) and our judge decoding when GPT-4o is drafting and L. ama - -00B verifying. Green denotes accepted and red rejected tokens. Right: Number of accepted tokens for GPT-4o as draft and L.1 a.lama-4 -4:0--4000 40B as target for standard speculative and our judge verification.
schemes Intuitively, one would expect from well-calibrated verification scheme to allow for accepting candidate tokens whenever they are contextually correct. However, as we show in the following two paragraphs, this not the case for standard logits based verification.
High-quality draft model. To further demonstrate how valid responses incur high rejection rates, we perform the following experiment: We take a very powerful LLM as the draft model and evaluate whether the target model accepts more candidate tokens, which are now guaranteed quality, While such a setup does not make sense for SD from an efficiency point of view (a powerful drafter is of course too slow), further investigates if acceptance rates improve with the quality of responses. To that end we use GFT-4o as draft model for the target Llama-405B. We generate full answers with Go on MT-Bench, GSM8K and HumanEval and simply check how many tokens the target accepts under greedy decoding before the first rejection, as there is way to properly perform SD) with closed-source models. order ensure that the target model able to "recognize" the high-quality tokens, we use the performant L1 Llama- 405B. We display the average acceptance length and an example prompt in Fig. 3 Counter-intuitively, we find that the target model does not reward the higher quality of tokens, accepting only roughly two before encountering the first rejection. To further explore if this observation changes when running the complete process of SD, reverse the roles of our standard setup and Llama 4 40BB draft for Ll amama 8B target model. Similarly, we find that reversing the roles reduces the number of accepted tokens slightly (see Fig. 2 right side), even though they are of better quality now. We thus conclude that acceptance rates do not improve with the quality of the responses.
be of high
Human expert drafting. Finally, we evaluate the efficacy of human annotations can- didate tokens for Llama-405B by processing Wikipedia articles. Using subset of the wikipedia summaary dataset (Scheepers) 2017). which contains high-quality, community- reviewed abstracts, we assess token acceptance rates under greedy SD verification when prompting the model to summarize these articles. As illustrated in Fig. 2 (right), substantial proportion of tokens face rejection, even within this high-quality context.
In summary, conclude that SD verification in its current form is highly inefficient, large por- tions of correct answers are rejected. Motivated by this insight, the following section presents more effective verification scheme that goes beyond model alignment in order to increase efficiency.
Additional examples of this phenomenon are provided Appendix C.I for further examination.
5

---

Published as a conference paper at ICLR 2025
Correct Answe::
it the deyina penssan atfacted by number chines, it world seill take houss to day 30 alsicka, just it took houre dry shieta.
Luresboos,, beltzenland, Lialy, Spains Anrdssiw, and Konaco.
Question:
Seput Prance?
Ie hang abdete octesde takes ther. hours dry, long
Krong Answer:
Prance shares St' borders wath Beiqian. Lccembeuurr, dereang,
weong Anave::
is would take 30 hours dry 30 chirte, an each shiet seeds an additionaa hour day.
Italy, apain, Rortogal, Poland.
Figure 4: Two examples from our dataset. We highlight the incorrect tokens in the wrong answer in red.
4 JUDGE DECODING
As demonstrated by previous experiments, need a more flexible method of verifying sequences to increase the number accepted draft tokens, especially as draft models continue to improve in quality. Recent work by [Zheng al. (2023) showed that large LLMs can reliably act as judges to evaluate responses generated by less capable models, correlating highly with human ratings. This judging approach allows for more versatile evaluation, focusing correctness and contextual qual- ity rather than strict alignment. However, using L.LM-judges directly is not feasible because (a) they require lengthy system prompts and often chain-of-thought reasoning, slowing inference, and (b) they evaluate full answers, whereas SD requires evaluating short, sometime partial continuations.
We thus aim to achieve this judge-like behavior efficiently while retaining the advantages of the orig- inal verification method, which ensures accurate next-token predictions in case of rejection. Since this involves computing embeddings for each draft token, we explore whether these embeddings contain sufficient information to enable rapid, reliable judgments.
4.1 VERSATILE AND ACCURATE VERIFICATION WITH TOKEN EMBEDDINGS
Token embeddings signal errors. Contrary to standard SD, which accepts or rejects a given to- ken based on its softmax probabilities (see Eq. 4 we find that the model's 's reaction processing the incorrect token itself reveals surprisingly valuable information. Specifically, our experiments show that last hidden layer embeddings of erroneous tokens effectively "flag" errors and contradic- tions, prompting the model to generate subsequent tokens that attempt to correct the mistake. This phenomenon is strikingly illustrated in Fig. where we condition L Llama- 40 5B on wrong replics (highlighted in red) and observe the model's immediate efforts to rectify its response (highlighted in green). For instance, when forced to start with the incorrect statement "The capital of France is Berlin", the model continues with "... just kidding, it's actually Paris". More such examples can be found in Appendix CC33 This unexpected behavior suggests the feasibility of leveraging the em- bedding of the current token following that a simple logistic regression head on top of these embeddings achieves high accuracy and can be trained in under 1.5 hours. Prior, embeddings have also been used to discover latent knowledge in LLMs and edit them (Burns et al. 2023 Zou et al. 2023a von Rütte et al.) 2024). further underscoring their richness.
means of verifying its correctness. In fact, will show in the
Dataset curation. In order to leverage token embeddings for verification, we carefully craft dataset consisting of high-quality user inputs, along with a correct and wrong answer pair. The set of input prompts are a mixture of newly-created questions and two public datasets that we heavily filtered (Alpaca (Taori et al.) 2023) and ARC (Clark et al. 2018)). Importantly, input questions and none of the answers. We leverage several models to produce a diverse set of cor- rect and wrong answers, including Mistral-Laree 2, LLama-82 and L amaa-055., thereby fostering robustness of the trained judge to recognize correct but differently aligned solutions. All answers were manually reviewed and corrected by the authors, who also annotated the precise lo-
only use the

---

Published as a conference paper at ICLR 2025
cation of errors in wrong answers In total we collected 500 high-quality question, correct answer, wrong answer tuples. For training, we label every token from the correct answer as positive, every token from the wrong answer up until the point of mistake positive, and finally every mistaken token as negative. Two examples from our dataset are depicted in Fig. Naturally, the dataset exhibits strong imbalance, leading to roughly 20x more positive than negative examples.
Model design and training. Equipped with the dataset, we train linear head fjudge top of the target embeddings, using a weighted cross entropy loss as the objective to counter the imbal- ance in the dataset. We place larger weight on the negative examples in order to ensure that the resulting judge does not falsely accept wrong tokens to limit quality degradation and perform early- stopping to reduce overfitting. We tune all hyperparameters with embeddings from several layers and find that deeper layers perform best with only insignificant differences, while shallow layers are clearly worse, consistent with similar observations in pre- vious works (Zou et al., 2023b] Gurnee & Tegmark, 2023] von Rütte et al., (024).. For simplicity, we thus stick to using the last embedding of the target before the RMS normalization (Zhang & Sennrich 2019) and the language modelling (I.M) head. While we experimented with more com- plex architectures, including MLPs and shallow Transformer networks, simple linear head proved most effective, demonstrating excellent performance without overfitting. This linear classifier of- fers significant practical advantages: we train only 16.4k parameters on just 30k tokens in less than 1.5 hours, with all target model parameters remaining frozen. Additional details are provided in Appendix B.I
small test split. We experiment
Inference. How the judge head now combined with the standard elements of SD? In essence, we use Judgee as an additional evaluator for given token (or rather its embedding e; E R and accept if a (Judge(e:)) > å for [0.5, 1], where a sigmoid function. In other words, serves as threshold for the confidence of acceptance and practitioners can thus choose how much to trust the judge layer. In practice we observed that there is no need to tune this value to ensure quality and leave it at the natural value 8 = 0.5. Given sequence of candidate tokens Ci....,, CM. thus get two accept/reject masks from the target model: Zstand E {0, }.M as in standard SD verification and Zjudge (0, 1 from the judge head. We take the logical OR between the two, 2 Zstand V Zjudge+ since when the judge rejects and standard SD accepts, the corrected token according to the target will exactly be the same token. We thus already accept it. We illustrate this mechanism in more detail in Appendix B.2 Note that d = reduces to SD.
Fs
PM--1
User Question:
User Question:
What is the capital of Prance? Tell something about city.
khat is 402 335?
Assistant:
Assistant:
Transformer Blocks
The capital France Serlin... No juat kidding. The capital of France is actually Paris. [...]
402 335 736 737
8x
CM
Figure 5: Left: Conditioning 1.l.am-4455 on wrong outputs. The part of the assistant response in red was forced, while parts in green were generated freely. Right: Judge illustration where s, is the last token from the context s and ...... CM are candidate tokens. Orange denotes embeddings, green denotes the LM-head output and red denotes the produced judgements.
Using L.LMs to that end proved he too imprecise, which is consistent with recent observations in Tyen 2024).

---

Published as a conference paper at ICLR 2025
Table Average acceptance length (m..) and speedup factor over standard decoding in HuggingFace and gpt-fast for batch size 1. We report generation tokens/s for gpt-fast for 512 imput and output tokens, quantized 8-bit. All 70B (405B) models run on (8) H100 GPUs, except for "Medusa Nvidia 2024), which runs on significantly faster H200s with NVLink Switch and TensorRT.
HUSSTNGFACE
GPT-FAST
TOKENS/S (512 + 512)
88/70R-STANTARS
6.4
1.5x
1.7x
76.7
93/70B-JUDGE (OURS)
18.8
2x
3x
141.8
RAG-RA
4.5
3.3x
1.9x
88.1
13/405B-STANDARE
6.3
5.3x
1.78x
58.7
83/405B- JUDGE (OURS)
19.7
9.7x
3.9X
129.3
4053-MEDUSA
6
6x
1.9x
108"
5 EVALUATION OF JUDGE DECODING
First, revisit the initial experiments outlined in Sec. 3.2 where use 4o the draft model, as well as human generated text as candidate tokens. both cases, the average number of accepted tokens is significantly higher for our method across datasets (see Fig and The example prompt in Fig. 33 (left) shows that the correct response of GPT-4o is fully accepted by judge decoding while standard SD rejects after words. This illustrates that verification scheme offers more versatile decisions.
5.1 PERFORMANCE BENCHMARK
We now evaluate our verification method on standard benchmarks in the SD literature. including GSM8K (Cobbe et al. 202T), HumanEval (Chen et al. 202T) and MT-Bench (Zheng al. 2023). In contrast to standard SD works, we do need to report the achieved accuracy values of our strategy, as adapting verification comes with the possibility formance. To give more complete picture, we further include multiple-choice benchmarks ARC Clark et al. (018)) and MMLU (Hendrycks et al. 2021), which are atypical tasks for standard SD as only a few tokens need to be produced, but further serves check that our verification scheme does not degrade performance. We use the prompting templates from Dubey et al. (2024).
accepting wrong tokens and thus worse per-
Training-free baseline. To provide more context for our results and to demonstrate that our judg- ing strategy goes beyond simple heuristics, also explore simple training- free method to relax the acceptance scheme. In particular, we investigate top-K K verification, where a candidate token e; accepted, if it is among the K highest valued probabilities p: produced by token G-11 1- K E N is a hyper-parameter of the decoding technique that trades-off quality against speed. Setting = V reduces running just the fast draft model, while K results in standard SD.
Preserving target performance. We display the accuracy of judge decoding alongside the vanilla draft and target models, as well as top-K decoding in Fig. (6 for both L. amaa4 and Llama-705 We observe that judge decoding almost exactly preserves target performance for all benchmarks, showing hence that up to 20 tokens can be accepted on average from modern draft models without loss of quality. The simple heuristic baseline, on the other hand, is hardly able to improve over the draft model (even for K 5), highlighting the difficulty of the problem we address with the learned head.
5.2 SPEED BENCHMARK
The end-to-end speed-ups achieved by SD methods improve as mainly two factors increase: (1) the number of accepted tokens and (2) the latency gap between draft and target model. Importantly, the latter is heavily dependent on whether or not orthogonal inference time optimizations like quantiza- tion, model parallelism and graph/kernel optimization techniques (like torch. torch.compile and TensorRT)
8

---

Published as a conference paper at ICLR 2025
Llama-8b
Blarna8b/70b--to 8b/70b-top--
Liama-8b/ saarr//////bbbuuuu
Llama-700
Accuracy (%)
100
91.
90
20-7
22.6
01.2 20.5
60
Llama-8b
Llama-Bb Bww--Bbb/--00-B 405b-top-k
Hlama-bb405bbjuddg
Llama-405b
100
26.5
96.8
20.6 96.6
Accuracy
90
25.2
20
86.7
12.0
19.7 50
80
71.3
GSMBK
HumanEval
ARC-Challenge
MMLU
MT-Bench
Figure 6: Benchmark results. Top: Draft Llama- 8B and target -1ama-700 -0B. Bottom: Draft Llama-8B and target Llama- 405B. We show top-K decoding, standard SD for M 10 and our judge decoding for .M (striped). Notice that our judging method preserves accuracies very well, while top- K loses most performance.
are applied. Unfortunately, prior works on SD have almost excursively relied on the ari-ffrienlly but un-optimized library HuggingFace (Wolf et al. 2020) to implement their methods. Yet, as rightfully pointed out by (Wertheimer et al., 2024). speed-ups of prior SD methods reported in HuggingFace tend to shrink significantly when moving to optimized inference frameworks. For ex- ample, the acceleration of Eagle on Liama- -amm--222 2 7B reduces from 3x to merely 1.5x when using gpt-fast reported in Li (2024a). In fact, vanilla L amma708 without any speculative decod- ing achieves a higher throughput in gpt-fast than the state-of-thh-aa SD method Eagle-2 does on the same model in HuggingFace (~455 vs <33 tokens/s).
To offer a complete picture, we here provide latency benchmarks in both frameworks, HuggingFace facilitate comparison, well as the arguably more relevant and optimized gpt-fast framework (Pytorch- Team 2023). If not stated otherwise, we run Llama- -0mm-70 and alamm -4 45B on and Nvidia HI00 GPUs respectively. Our results are summarized in Table
Llama-3.1-70B. When drafting with Llama-88 8B for Glamm- -0m with batch size 1 in simple frameworks, the latency delta between the two models is relatively small, limiting the speed-ups of judge decoding. This is particularly evident when compared to SD methods that leverage small LM heads as draft modules (such as Eagle-2 Li et al.) (2024b) and Medusa Cai et al. (2024)). However, in the several latency bottlenecks (like CPU instruction and memory 1/O) are alleviated, resulting in more pronounced latency delta between the target and draft models. Consequently, our method effectively capitalizes on this increased latency disparity and outperforms the current state-of-the- art by substantial margin (see right-hand side of Table T.
realistic setting of deployment within optimized inference framework,
Llama-3.1-405B. Replacing the target model with the powerful ..laa-4055 model sig- nificantly increases verification latency. As result, drafting (and accepting) longer sequences be- comes more crucial for the overall runtime. In such settings, judge decoding shines because the average number of accepted tokens is > 3x larger than prior works (left-hand side of Table In particular, both Medusa and Eagle: e --ar limited drafting & 6 token the time, by the num- ber of heads and the draft tree depth respectively. Our 8B/ 405B-Judgee however, accepts close 20 tokens at time and thereby achieves 9.7x speed-up HuggingFace and unprecedented 129 tokens/s in gpt-fast.
5.3 OUT-OF-DISTRIBUTTIN PERFORMANCE
Finally, we investigate to what degree our judge-decoding strategy extends to situations for which has not been trained. To this end, we filter our dataset by removing all coding examples, train the verification head for Llama-B8 405B on this reduced set and then evaluate on the coding task
of which 405B version exists.

---

Published as a conference paper at ICLR 2025
HumanEval. While we do observe a drop in performance from 86.6 to 80.4%, the performance is still significantly better than the draft model at 71.3%, indicating that the notion of "correctness" transfers between tasks at least to some degree. Nevertheless, approach is not silver bullet and to maintain target quality it is required to train the judge on data of similar nature.
6 CONCLUSION
In this work we have investigated the verification mechanism in speculative decoding and identified how its focus on alignment between draft and target response leads to the rejection of objectively correct continuations. To fully leverage the improved quality of "small" language models, we thus proposed an adapted verification scheme that makes use of the capability of LLMs to judge responses in versatile way. This allows for efficiently drafting more tokens, leading to significant speedups up to x on a range of benchmarks, achieving unprecedented speeds of 129 tokens/s for a.lama-400 -0.B. In the regime of many draft tokens, "small" language models shine drafters compared to the small modules employed in approaches like Eagle e or Medusa and we believe this trend will only further accentuate in the future. Our approach however also comes with drawback; the mathematical guarantee to maintain target quality is lost by relying on the judge. Through extensive experiments we show that a well-trained judge does not lose performance on standard benchmarks and we thus view our approach as a significant first step into this direction. On the other hand, our strategy in its current version does not present silver bullet; novel tasks require the careful annotation of similar data to maintain quality, otherwise performance is lost. The small amount of data required in our setup is nevertheless very encouraging sign. Future work can hopefully build upon our contributions, further improving our judge decoding strategy to enable more speedups.
REFERENCES
Marah Abdin, Jyoti Aneja, Hany Awadalla, Ahmed Awadallah, Ammar Ahmad Awan, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Jianmin Bao, Harkirat Behl, Alon Benhaim, Misha Bilenko, Johan Bjorck, Sébastien Bubeck, Martin Cai, Qin Cai, Vishrav Chaudhary, Dong Chen, Dong- dong Chen, Weizhu Chen, Yen-Chun Chen, Yi-Ling Chen, Hao Cheng, Parul Chopra, Xiyang Dai, Matthew Dixon, Ronen Eldan, Victor Fragoso, Jianfeng Gao, Mei Gao, Min Gao, Amit Garg, Allie Del Giorno, Abhishek Goswami, Suriya Gunasekar, Emman Haider, Junheng Hao, Russell J. Hewett, Wenxiang Hu, Jamie Huynh, Dan Iter, Sam Ade Jacobs, Mojan Javaheripi, Xin Jin, Nikos Karampatziakis, Piero Kauffmann, Mahoud Khademi, Dongwoo Kim, Young Jin Kim, Lev Kurilenko, James R. Lee, Yin Tat Lee, Yuanzhi Li, Yunsheng Li, Chen Liang, Lars Liden, Xihui Lin, Zeqi Lin, Ce Liu, Liyuan Liu, Mengehen Liu, Weishung Liu, Xiaodong Liu, Chong Luo, Piyush Madan, Ali Mahmoudzadeh, David Majercak, Matt Mazzola, Caio César Teodoro Mendes, Arindam Mitra, Hardik Modi, Anh Nguyen, Brandon Norick, Barun Patra, Daniel Perez- Becker, Thomas Portet, Reid Pryzant, Heyang Qin, Marko Radmilac, Liliang Ren, Gustavo de Rosa, Corby Rosset, Sambudha Roy, Olatunji Ruwase Olli Saarikivi, Amin Saied, Adil Salim, Michael Santacroce, Shital Shah, Ning Shang, Hiteshi Sharma, Yelong Shen, Swadheen Shukla, Xia Song, Masahiro Tanaka, Andrea Tupini, Praneetha Vaddamanu, Chunyu Wang. Guanhua Wang, Lijuan Wang. Shuohang Wang, Xin Wang. Yu Wang, Rachel Ward, Wen Wen, Philipp Witte, Haiping Wu, Xiaoxia Wu, Michael Wyatt, Bin Xiao, Can Xu, Jiahang Xu, Weijian Xu, Ji- long Xue, Sonali Yadav, Fan Yang, Jianwei Yang, Yifan Yang, Ziyi Yang. Donghan Yu, Lu Yuan, Chenruidong Zhang, Cyril Zhang, Jianwen Zhang, Li Lyna Zhang. Yi Zhang, Yue Zhang, Yunan Zhang, and Xiren Zhou. Phi-3 technical report: A highly capable language model locally on your phone, 2024. URL
arbb abs/ 2404. 1.e..145. 14219
Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Ragan- Kelley, and William Brandon. Hydra: Sequentially draft heads for medusa decoding, 2024. URL htt.:ttpss//aai..o////a/aa//////
05109
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Forty first International Conference on Machine Learning, 2024. URL https: http::tttt////eeeet httpp::s//ppeeeviee.. net/ forum? d= 67.
10

---

Published as a conference paper at ICLR 2025
Nikhil Bhendawade, Irina Belousova, Qichen Fu, Henry Mason, Mohammad Rastegari, and Mahyar Najbbi Speculative streaming: Fast Ilm inference without auxiliary models, 2024. URL hetps:
//axxiv.rg/abbs 2ttpa//2222.111
11131
Collin Burns, Haotian Ye, Dan Klein, and Jacob Steinhardt. Discovering latent knowledge in lan- guage models without supervision. In The Eleventh International Conference on Learning Rep- resentations, 2023. /
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv: 2401. 10774, 2024.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean- Beann-Batiiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling, 2023. URL
https:///axixs org/ abs 2302.01318
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katic Mayer, Peter Welinder, Bob Mc- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code, 2021. URL http:: /ttp:::///rriiii org/ abs /tt.b//1111 22107 21077033777
Zhuoming Chen, Avner May, Ruslan Svirschevski, Yuhsun Huang, Max Ryabinin, Zhihao Jia, and Beidi Chen. Sequoia: Scalable, robust, and harawere-aware speculative decoding, 2024. URL
12374
Jacob K Christopher, Brian Bartoldson, Bhavya Kailkhura, and Ferdinando Fioretto. Speculative
diffusion decoding: Accelerating language generation through diffusion, 2024. URL https:
//axxiv.ogg/as//222... 05636
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, abs/1803.05457, 2018. URL stttppppsepppp...mmiiiiiiieeee .cscholrroooogg
CorpusID::392 3922811
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano. Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https:
org/ bog//bs//a 2110. 14168
Cunxiao Du, Jing Jiang, Xu Yuanchen, Jiawei Wu, Sicheng Yu, Yonggi Li, Shenggui Li, Kai Xu, Liqiang Nie, Zhaopeng Tu, and Yang You. Glide with caPE; A low-hassle method to accelerate speculative decoding. In Forty-first International Conference Machine Learning, 2024. URL htttps://openrevs eww net. larte/oorune forum? 10=
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song. Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Ander- son, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah
11

---

Published as a conference paper at ICLR 2025
Korevaar, Hu Xu, Hugo Touvron, Ilyann Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Ma- hadcokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Al- wala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala- Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Man- nat Singh, Manohar Paluri, Marcin Kardas, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhar- gava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sum- baly. Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa. Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang. Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao. Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vladan Petro- vic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen. Zoe Papakipos, Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisemman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti. Bran- don Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Ar- caute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan. Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Gold- man, Ibrahim Damlaj, Igor Molybog, Igor Tufanov. trina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang. Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Mathhia Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Ke- neally, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mo- hammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navy- ata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong,
12

---

Published as a conference paper at ICLR 2025
Ning Zhang, Norman Cheng, Oleg Chemoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittnrr, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Rohan Maheswari, Russ Howes, Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Sa- tadru Pan. Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lind- say, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shiva Shankar, Shugiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max. Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Tim- othy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vitor Albiero. Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang. Wenwen Jiang, Wes Bouaziz, Will Con- stable, Xiaocheng Tang, Xiaofang Wang, Xiaojian Wu, Xiaolan Wang, Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen. Zhenyu Yang, and Zhiwei Zhao. The llama 3 herd of models, 2024. URL https: /ttpsv//aaaiv... org/ abs /200 21783
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed Aly. Beidi Chen, and Carole-Jean Wu. LayerSkip: Enabling early exit inference and self-speculative decoding. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 12622-12642,
Bangkok, Thailand, August 2024. Association for Computational Linguistics. URL https:
//aclanthology org. 22044 acl long. 681
Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of LLM in- ference using lookahead decoding. In Forty-first International Conference Machine Learning, 2024. htattts::/openrevviie http://www...ett.oooooooo/aaaaaa
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Roziere, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction. In Forty-first International Conference on Machine Learning, 2024. URL httph:tt/////t.. /opp://v///peeevviii ew net. forum? d= pEWAceji0D
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. 752
URL (ttL:ttp::// http:////xxvvvo/eeeeeee
Wes Gurnee and Max Tegmark. Language models represent space and time. arXiv preprint arXiv:2310.02227,, 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Ja- cob Steinhardt. Measuring massive multitask language understanding. In International Confer- ence on Learning Representations, 2021. URL htttps://ppezeews net.eeeweeeeeeeeeeeeeeeeeeeeeeee forum?id
d7KB jm 3GmQ.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL
2203.
Kaixuan Huang, Xudong Guo, and Mengdi Wang. Speedec++: Boosting speculative decoding via adaptive candidate lengths, 2024. URL
/ttbs///25.. 19715
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems, 3.711-732, 2021.
13

---

Published as a conference paper at ICLR 2025
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. http:://: http:///axvvvv.. www.w.twww......a.a.oooooo.orete http://.s///00...oom 08361
Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W. Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. In Thirly-seventh Conference on Neural Information Processing Systems, 2023.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning, (C. JMLR.org, 2023.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty. In International Conference on Machine Learning, 2024a.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-2: Faster inference of language models with dynamic draft trees, 2024b. URL httLl/tttps:/arx....oorr
Fangeheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, and Yunhe Wang. Kangaroo: Lossless self-speculative decoding via double early exiting, 2024a. URL https:///axxiii
org/ arg/abs// 2404. 18911
Jiahao Liu, Qifan Wang. Jingang Wang, and Xunliang Cai. Speculative decoding via early-exiting for faster LLM inference with Thompson sampling control mechanism. In Lun-Wei Ku, An-
dre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Lin- guistics ACI. 2024, pp. 3027-3043, Bangkok, Thailand and virtual meeting, August 2024b. Association for Computational Linguistics. URL
Eindings ac1.179
Tianyu Liu, Yun Li, Qitan Lv. Kai Liu, Jianchen Zhu, and Winston Hu. Parallel speculative decoding with adaptive draft length, 2024c. URL (Rtppts:://aaxx
wttp 11850
Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding. arXiv preprint urXiv:2310.07117, 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2019. URL htttss:/ nwt.:nee//eeee... http://wwwwweeeeoouueeumm
Bkg6RicqY7
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating large language model serv- ing with tree-based speculative inference and verification. In Proceedings of the 29th ACM International Conference on Architecturul Support for Programming Languages and Operat- ing Systems, Volume 3. ASPLOS '24, pp. 332-949, New York, NY. USA, 2024. Association for Computing Machinery. ISBN 9798400703867.
doi: 10.1145/3600666.35511
URL
patpps:/doo..om
1111111153366661 3651333
Giovanni Monea, Armand Joulin, and Edouard Grave Pass: Parallel speculative sampling, 2023.
URL httttp:://axii...tot /as///b///// /ats//211... 2311.
Nvidia. Llama 3.1 performance with medusa. fttps::/sshotur/// att IIXXCC 2024. Ac-
cessed: 2024-10-01.
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Moham- mad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brock- man, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey
14

---

Published as a conference paper at ICLR 2025
Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus. Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gib- son, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hal- lacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey. Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Lukasz Kaiser, Ali Ka- mali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung. Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning. Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reichhro Nakano, Ra- jeev Nayak, Arvind Neelakantan, Richard Ngo. Hyeonwoo Noh, Long Ouyang, Cullen O' Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross. Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Sel- sam, Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalic Staudacher. Felipe Petroski Such, Natalie Summers. Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Pre- ston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vi- jayvergiya, Chelsea Voss. Carroll Wainwright, Justin Jay Wang, Alvin Wang. Ben Wang. Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wo- jeiech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report, 2024.
Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5:606 2023.
Pytorch-Team. gpt-fast. https://giihbbb.cm//ptttrrct hatpr/torrchhhhbs/aabe labs /pt fast 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Aishwarya P S. Pranav Ajit Nair, Yashas Samaga L, Toby James Boyd, Sanjiv Kumar, Prateck Jain, and Praneeth Netrapalli. Tandem transformers for inference efficient LLMs. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 42906 42917. PMLR, 21-27 Jul 2024. URL https::/proceee:: ngs. www.sr....s........ wt3.//2355///2
html
Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorea, Michele Mancusi, Ric- cardo Marin, and Emanuele Rodola. Accelerating transformer inference for translation via par- allel decoding. In Anna Rogers, Jordan Goyd-Graber and Naoaki Okazaki (eds.), Proceedings of the l1 Annual Meeting of the Association for Computational Linguistics (Volume I: Long
15

---

Published as a conference paper at ICLR 2025
Papers), pp. 12336-12355 Toronto, Canada, July 2023. Association for Computational Linguis-
tics. doi: 10. (0.188533v112023... aac.23...c-ooo.8..... URL attpp:://aclanttaoooo..ooo
acl- acl-long.689
Thijs Scheepers.
Improving the compositionality of word embeddings. Master's thesis, Universiteit van Amsterdam, Science Park 904, Amsterdam, Netherlands, 11 2017.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL httpps://ppeeeeeieeeeeeeee/eeeee net forum?
Noam Shazeer. Fast transformer decoding: One write-head all you need, 2019. URL https:
//araiv.orr//ab//111111111 02150
Mitchell Stern, Noam M. Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. In Neural Information Processing Systems, 2018. htttt::t///aaa
sseanttischoo
Xin Sun, Tao Ge, Furu Wei, and Houfeng Wang- Instantaneous grammatical error correction with shallow aggressive decoding. In Chengqing Zong. Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 5937-5947, Online, August 2021, Association for Computational Linguistics. doi: 10.18653/ v1/2021 1al.0211.c--ooo.440
/Btp:/tpss//cllaahoooooo.oog/// wtt.ooo///22...lll..ooo.. acl-long.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An foolloiion-foloowing Ilama model. https: / com htt///ttsuu-ee labulllabaae habbsstanford.com alpaca 2023.
Gladys Tyen. Hassan Mansoor, Victor Cärbune, Peter Chen, and Tony Mak. Lims cannot find reasoning errors, but can correct them given the error location, 2024. URL htttps://arxi>
brg//bb/2311 23111. 68516
Dimitri von Rütte, Sotiris Anagnostidis, Gregor Bachmann, and Thomas Hofmann. A language model's guide through latent space. arXiv preprint 20..2332..1433,, 2024.
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models, 2024. URL https:// /tttpp://arxi..../ org/ abs/
2448.15233
Davis Wertheimer, Joshua Rosenkranz, Thomas Parnell, Sahil Suneja, Pavithra Ranganathan, Raghu Ganti, and Mudhakar Srivatsa. Accelerating production iims with combined token/embedding speculators, 2024.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-thh-arn natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 48, Online, October 2020. As- sociation for Computational Linguistics. URL httph:://///.... www. aclweb. www.//raaogggggooooooooo
2020. emnle- demos.e
Heming Xia, Tao Ge. Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 3909-3925, Singapore, December 2023. Association for Computational Linguistics. doi: 111865331122223 findings-emmlp.2577 Findings emnlp.257
URL http::::s//aaaaahlllogeo org.org/////. 2023.
Aonan Zhang, Chong Wang, Yi Wang. Xuanyu Zhang, and Yunfei Cheng. Recurrent drafter for fast speculative decoding in large language models. arXiv preprint arXiv:2403.09911 2024a.
16

---

Published as a conference paper at ICLR 2025
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https://
waxiv..or//ao wtt.:rgggabb/..iiio 07467
Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. Draft verify: Lossless large language model acceleration via self-speculative decoding. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meet- ing of the Association for Computational Linguistics (Volume Long Papers). pp. 11263- 11282. Bangkok, Thailand, August 2024b. Association for Computational Linguistics. URL attts://accaane wow.ohoooo....oooo ogy..og......................... www.o////222.. http:/222...cll1l111
607
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang. Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging LLM-as-al judgee with MT-bench and chatbot arena. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023. URL https:
//openreview.netto net./ dww.ccpPPccccccccccccc.c.
Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean- Fran-François Kagy, and Rishabh Agarwal. Distillspecc: Improving speculative decoding via knowledge distillation. In The Twelfth International Conference on Learning Rep- resentations, 2024. URL wttpt:tt:///tmm..t. http:::://ppeervvleeee
net/ forum??GGrSY663377T
Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo. Richard Ren, Alexander Pan. Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, Shashwat Goel, Nathaniel Li, Michael J. Byun, Zifan Wang, Alex Troy Mallen, Steven Basart, Sanmi Koyejo, Dawn Song, Matt Fredrikson, Zico Kolter, and Dan Hendrycks. Representation engineering: A top- down approach to ai transparency. ArXiv, abs/2310.01405, 2023a. URL hhttps://apl
semanticscholar. ott.:poc......ooooooooooaaaaaaa.
Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, et al. Representation engineering: A
top-down approach transparency. arXiv preprint arXi::2310.01400 2023b.
17

---

Published as a conference paper at ICLR 2025
A LIMITATIONS
Here we list limitations of our approach to the best of our knowledge:
An obvious limitation is the loss of the mathematical guarantee to match target quality. While we perform extensive experiments and show that quality is maintained, there is no certainty for novel tasks.
The draft model needs be of high-quality, otherwise our approach naturally does not prove beneficial and too many tokens end up being rejected. self-speculation and small drafters the spirit of Medusa or Eagle are thus not ideal since their generations quickly deteriorate when drafting too far into the future.
Similarly, the target model needs to be of sufficient size to be able to provide accurate judge- ments. Speedups for smaller models such as Llama-88 8B are hence tougher to achieve.
As highlighted in the main text, new tasks do require careful annotation of data maintain quality. The required amount on the other hand turns out to be small in our case.
If the draft model has safety issues, the target model could potentially accept safety-critical tokens through the judge, even if the target would otherwise never produce such outputs. We have not observed such issues in our experiments but have also not thoroughly investi- gated this problem as it is beyond the scope of work.
B ARCHITECTURAL AND EXPERIMENTAL DETAILS
B. LINEAR HEAD
We train our linear heads using the AdamN optimizer (Loshchilov & Hutter 2019) with learning rate y= 0.0001, weight decay 0.1 and batch size 128. Note that for L1ama- 405B, our linear head has dimension 16, 384 while for ama ama-708 it has 8, 192. Our linear head can be viewed as an additional entry the vocabulary V. reducing its inference overhead thus to practically zero
B.2 JUDGE MASKING
We describe the the combination of standard and judge mask in more detail in Fig. In the following
Judge mask
Hazzy
Pottez
sexies
o1
iantaay
thsillers
wr written
Standard
Harry
Potter
eerles
01 tantaay
thsillers
wri written
mask
Final mask
Farry
Potter
serion
Fartasy
thri 1lers
Figure 7: Illusraaioon of mask creation in judge decoding. The decision mask resulting from the
judge is combined with the standard mask from SD. Once both methods disagree, subsequent tokens get rejected automatically as usual, even if they were individually accepted.
we will describe in more detail why it is natural to combine the masks of judge decoding and standard speculative decoding. We illustrate this in Fig. 8 There are (rare) scenarios where a candidate token is rejected by judge decoding (such as "that" in the example) and the "corrected" token according to the target model happens to the same token ("that" in blue). This situation could repeat; the very next token could again be rejected by JD and accepted by the target (token "guy" in the example). Standard SD on the other hand would accept all those tokens as the draft exactly matches the target suggestion. We eventually end up accepting the exact same tokens (in case of rejection we have to trust the target), so it makes sense to combine the masks and use the SD mask to not end up repeating the steps (in the example we do not observe this situation too often, but can occasionally occur as the judge was tuned to rather reject than accept when in doubt avoid false positives.
combine steps 1., 2., 3. into one step the right). In experiments
18

---

Published as a conference paper at ICLR 2025
contest
Accepted
Repected
Corrected
lokan
Phal
Short
Phaf
JD Mask
ever
that
Iouched
JD Mesk
Short
that
Icuehrod
SD Mask
Mask
Combined
Shortt
evaryytheng
that
Icustred
Mask
3.
Short
everything
Ihaa
loucheed
JD Mask
Figure 8: Illustration how combining the masks of judge decoding (JD) and standard speculative decoding (SD) results in the same reply but in less steps in the (rare) case that JD rejects a token that actually the target token.
B.3 HARDWARE
We run all of our experiments on a single node of H100-SXM5 GPUs. For 1.lam--405B we use 8 GPUs and 8-bit quantization ensure that the model fits on a single node. For Ll ama 70B, use again 8-bit quantization but only 2 GPUs.
C MORE PROMPTS
C.I REJECTED REPLIES FOR LLAMA-88
Here we provide more example prompts where Llama- 8B provides completely correct answers but gets rejected early by the target Llama-8a 405B. We display the decision of judge decoding strategy right below. To also highlight that judge decoding can catch errors and does not just blindly accept responses, also show prompts where L. amlmma-8 8B provides wrong response.
4b/4036 Standard Necification:
tb/425b Standard Veridication: will salculate the wonthly share
Usut Quantion:
find out
vehicie
Dauz Queation:
bua tzeveia 6c ediea per hous for houra.
Camma bus nonth the uimuiricily bill in $100, Row nuch will each coormate pay per yeaz dieiide shace squally?
#b/40tb Zudge Verificatioon:
cind out wuch vaticie
Leaveled, nasd miltyppy
Zudge
bourz. auch
tarther did the
to cind out tox euch each roomsaee per years
bus yo chan the cas, ailes?
The bus teaveled 60 wiles hour for houes, so it. exeveled 60 300 nilen.
nuttiply aonialy
The moothly eleersectty ai1: ie $100. stnce thare are coomnatee, Me by
for heure, exsseled 30 240 eiles.
find nush darebar bes traveled Zeom dtatance Suz traveled.
Figure 9: More example prompts for SD for L. 1.lama-20 and Llama-405B. Left: Correct response getting rejected early under standard decoding, while judge decoding accepts long continuation (but admittedly over-cautiously rejects later on). Right: Wrong response that gets rejected too early by standard decoding and correctly rejected later on by judge decoding (there are roommates in total).
19

---

Published as a conference paper at ICLR 2025
9b/405b Standard Versfication:
8b/405b Standard Verification:
Dear Conation:
"Island of Alcha
Dear Gonation:
David
lavic sad thoee siataze. Each of than has ose nany
travel bloy post stout recent crip to Raws sighiigoting
*Islend of Alcha: Devulling the zidden Gaea of Havail
lu-b//425b Judge Verification:
atapped the plans
fwoid have?
Scnce bavld three sisters, and wuch
sunolaie,
Lrupicsi aut
beothee, Lial Bavid bzothar.
storactions.
pluseria filled
finally in Hovall the uitirate pacadiie on eazth. By recest trip to the tha State vas Journey discevasy, lamazaing acsait and unforgetable expertentes that
nake Hevall
vrpultural Expeziencepre
nelcing pot Tooel
Figure 10: More example prompts for speculative decoding for Llama-88 and 1lama-400 Left: Correct response getting rejected early under standard decoding, while judge decoding accepts long continuation (but admittedly over-cautiously rejects later on). Right: Wrong response that gets rejected too early by standard decoding and correctly rejected later on by judge decoding (there 5 roommates in total).
C.2 JUDGING OF WIKIPEDA ARTICLES
Here provide more details and examples for verifying Wikipedia articles. Given Wikipedia article name such as "aluminium", "Moore's s Law" or "Pet Shop Boys", we prompt the model for information by asking "What again compare greedy matching for standard speculative decoding with our judging strategy when using the summary part of the Wikipedia article reply. We display some example prompts along with the corresponding verifications in Fig.
you tell about <insert topic>?". We then
Oser Dreation:
Zizc chemioes aleeent with sysbo2
Dseation:
Doznie Sucta 2001 Amazican
Rlst can, you sall
What can you
Judge
sive ceneett with and stondo neabes 3t. 2o La the daret
nire file
by
Richard Bally, 25 abaza dase Opilantast Jaza Nalena, Dewn Daczynoce, Janan
Figure 11: Two example prompts from the subset wikipedia-an soniiiiaassmariiss..... along with the correspond verifications.
C.3 FORCING WRONG REPLIES FOR LLAMA - -405
Here we provide some more evidence of the "correcting" behaviour of 1.lama44052 when condi- tioned on wrong tokens. If the model cannot fix the response anymore, then will often point that the completion it just gave is actually wrong, see e.g. examples in Fig. [3] This again strongly suggests that correctness should thus be detectable the embeddings of such tokens.
20

---

Published as a conference paper at ICLR 2025
User Question:
User Question:
Write me Python function that doubles every element in giver list.
What 1s heavier: kilogram ot feathers or k-logram stones?
Assistant:
Assistant:
def double elements (1):
A kilogran of stones is heavler because... just kidding. They aze actually the seme weight.
"""Doubles every element in list for in range (1)
1[1] - 3 1/11 1[4]
return
Figure 12: Left: 2lama-400m corrects the mistake by subtracting 1 [i] to double instead of triple. Right: Similar correction behaviour by pointing out that response so far is wrong.
Question:
Question:
Write me Python Function that doubles every given 14st.
Bython function every entry in to zero and appends the result the origina: wanodirled list.
Assistant:
Assistant:
double :
add sooseroll)
***Doubles every elenent in listnes for i in (len(1)). 111]
aneSet every odd entry In 11at to sero and append the original unnodi listhrs copy
111]
return
for range (lenicopyl::
copy
However, the above function a bug. Here is the cocrested version [...]
copy(i) 0
return copy
Bowever, the above function will modicy the is varisble. dere is the corrected version [...]
original
Figure 13: Left: Llama-400 405B correctly points out that there is a bug as the index variable "j"" is not defined. Right: The model also catch more subtle mistakes. Here the original list also gets modified as no copy was made, leading to wrong outputs.
21


Reviews:
[{"review_id": "Csm6wLUKaK", "reviewer": "Reviewer_yijH", "summary": "The submission proposes Judge Decoding (JD), a variation of Speculative Decoding (SD) for LLMs. Like SD, JD uses a small draft model to propose output, and uses a larger model to check the output. The difference between SD and JD is in how the check is performed: where SD focuses on checking that the draft model's output is exactly like the larger model's (i.e., effectively per-token equality), JD uses a weaker criterion, based on the embeddings of the final layer of the larger (judge) model. The downside of this weaker criterion is that the output of the model is not guaranteed to be equivalent to the larger model's output anymore, but experiments indicate a significant potential speedup.", "questions": "* Fig 2/Fig 3: It would be useful to report the number of proposed tokens as well, either explicitly or by presenting \"ratio of accepted tokens\" rather than an absolute number.\n* One surprising result is not investigated: using Llama 405B to draft outputs of the 8B model on GSM8K (lower right of Fig. 2) seems to work less well than the other way around, even though the models are obviously well-aligned. I would have expected a basically equal result here (following the intuition that the larger model is effectively a superset of the smaller model), and was very surprised to see a lower score.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* Clear presentation, covering both technical details and high-level ideas\n* Comprehensive experiments that emphasise the strengths of the proposed method\n* Reported speedups seem significant, and seem like they could be very useful for practical serving of LLMs", "weaknesses": "* Some of the results (e.g., lack of agreement between different models / human drafts and models) seem to be obvious and do not help to make the argument\n* The use of a disjunctive criterion (standard SD accepts or JD accepts) is a bit surprising, and left me wondering in what cases (and how often) JD rejects when SD accepts. The paper would be strengthened by investigating this more closely.\n* Sect. 5.3 (OOD performance) is very important for deployment of the proposed method, and so it would be useful to expand on this. In particular, Sect. 4.1/\"Dataset Curation\" does not go into much detail into what distribution the training data follows, and which tasks are covered well. It would be helpful to understand whether there is a strong alignment between the TokenCourt training data and the considered evals.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FaWM45VIAp", "reviewer": "Reviewer_kfoq", "summary": "The paper introduces Judge Decoding, a speculative decoding enhancement that improves token acceptance by evaluating contextual correctness rather than strict alignment with the target model. Traditional speculative decoding often rejects accurate tokens from draft models due to misalignments rather than incorrectness, limiting speedups. By training a lightweight “judge” module on a custom TokenCourt dataset, Judge Decoding accepts more valid tokens, achieving up to 9× faster generation speeds with maintained quality on the Llama-3.1 model family, demonstrating a promising approach for efficient large language model inference.", "questions": "NA", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written. The argument is clear and reasonable. The proposed method shows effectiveness.", "weaknesses": "I do not see significant weakness.\n\nBut the sentence \"To speed up inference in such a setting, Chen et al. (2023); Leviathan et al. (2023) concurrently propose speculative decoding (SD)\" in the introduction section is incorrect.\nAt least, there is a much earlier one in 2018 that proposed speculative decoding: \"Blockwise Parallel Decoding for Deep Autoregressive Models\" by Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. The author skipped a large amount of literature on speculative decoding before 2023.\n\nAuthors are required to address this issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5OdoEug7Pm", "reviewer": "Reviewer_h6mH", "summary": "The authors propose \"Judge Decoding,\" a novel approach to speculative sampling that challenges the fundamental assumption that verification must be based on model alignment. Instead of using likelihood comparisons, they train a simple linear classifier on the last embedding layer of the target model to judge token quality directly. Their key insight is that current speculative decoding methods reject many high-quality tokens simply because they don't perfectly align with the target model's preferred phrasing. By analyzing embeddings and observing how models attempt to correct errors, they show that correctness information is already present in the token representations. They leverage this insight to create a lightweight judge module (16.4k parameters) trained on their carefully curated TokenCourt dataset. The method achieves up to 9× speedups over standard decoding, reaching unprecedented speeds of 129 tokens/s for Llama-405B, while maintaining quality across a comprehensive range of benchmarks.", "questions": "No questions", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Clever idea that was well-executed. The paper is well-written and the experiments make sense. The results are impactful in practical terms (people should use this) and for spawning future research (there are many possible variations on how to train the linear classifier).\n\n- I liked that the authors decided to go with the simple linear classifier since it worked best, rather than overcomplicating things\n\n- I also liked the systematic investigation with illustrative examples for how judge decoding differs from speculative decoding\n\n- Achieves state-of-the-art performance (141 tokens/s for 8B/70B-Judge, 129 tokens/s for 8B/405B) on standard hardware (2-8 H100 GPUs).", "weaknesses": "I would have liked if the authors addressed the possible impact of their method on less binary (correct vs incorrect) metrics like reward model scores or chatbot ELO scores. I could imagine to get a more noticeable regression on those, since the writing style of models is more important in that setting. Perhaps something for the discussion?\n\nThe Burns et al 2022 paper (https://arxiv.org/abs/2212.03827) does some nice work on eliciting concepts with linear probes which could be relevant here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ObW1XPx2u1", "reviewer": "Reviewer_Sf3Q", "summary": "In standard Speculative Decoding (SD), draft tokens are typically accepted or rejected during the verification phase based on whether they align with the target model distribution. However, this paper raises the concern that this verification process may be too stringent, resulting in high-quality draft tokens being rejected because they don’t align to the target distribution even though they may be valid and correct continuations. The authors try to address this concern by building upon the idea of LLM-as-a-judge and training a small judgment model that works on top of the target model latent embeddings and can be used to judge whether a draft is a correct continuation or not. Using this idea, the authors present speed ups using draft and target models in the Llama-3.1 family on a range of benchmarks and present a novel dataset (TokenCourt) used to train the target judge model.", "questions": "- Do the authors plan on releasing their TokenCourt dataset?\n- What happens when the target LLM isn’t one of the models used to generate the correct/incorrect responses to questions in TokenCourt used to train the linear judge module? Is this a necessary requirement for this technique to work? I ask because of a scenario where one may want to re-use a curated TokenCourt dataset for a new target model from an entirely different family of models to avoid having to hand annotate again.\n- Looking at the showcased TokenCourt examples, one question is if Judge Decoding may result in potentially verbose outputs. And in those cases, even though more draft tokens are being accepted and decoded per second and the full output is semantically correct, it is not as concise or clear as the original target model output if standard verification were being applied?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The authors clearly present and motivate their ideas around standard SD verification potentially being overly stringent and distinguishing between semantic correctness and target distribution “alignedness”. The motivation is straightforward to understand with the examples provided and it’s clear how standard verification can be arbitrarily strict depending on the target model resulting in less than desired inference efficiency using standard SD at times. \n- The reported speedups over standard SD are significant, largely in part due to the increase in the number of accepted tokens using this new judge decoding verification process compared to standard verification. However, the authors correctly point out also that this is not necessarily an apples to apples comparison given the loss of output distribution guarantees compared to standard SD.\n- The investigations into token embeddings signal errors was particularly interesting in that “incorrect” statements in text could be pinpointed based on the target model’s latent embedding because of the models tendency to try to immediately rectify its response following an inaccuracy and the resulting dataset curated based on this could prove useful to others if released", "weaknesses": "- Given the fact that verification is being changed so there is no longer a guarantee of output distribution match, the speedup comparisons are no longer apples to apples and have to also be paired with quality comparisons to ensure quality is maintained. Additionally, given the way the judge model is trained on “correct” and “incorrect” inputs, it seems that for novel tasks, a new tuning dataset must be hand curated for each task to train the judgment model to recognize “correctness” making the overhead of applying a technique like Judge Decoding significant compared to other SD modifications.\n- The notion of “correctness” is easiest to understand in the context of text generation tasks that are Question&Answer based. However, it’s unclear how this idea would work for more open-ended creative generation tasks where “correctness” may not be as easily defined and matching the target model output distribution may be ideal.\n- An additional condition to applying Judge Decoding is for the draft model to be of high-quality that can produce longer valid generations otherwise the approach may not provide that much of a speedup. \n- The reported experimental results are somewhat limited such as the out of distribution performance results being restricted to only a few lines and the full results not being reported in the appendix at least in comparison to the baselines etc. Additionally, it’s unclear what task is being reported on for Table 1 whereas Figure 6 shows accuracy across a variety of tasks. It would be helpful to also see the relative speedups across each of these variety of tasks as well.\n- For the most part, the authors are forthcoming about weaknesses and limitations to their judge decoding approach and many of these points will align. However, the amount of pre-requisites needed to apply the technique as well as required overhead for verification of maintained quality due to the loss of output distribution alignment are substantial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)