Speculative Sampling via Exponential Races (2025)

Abstract: Speculative decoding accelerates large language model inference using a smaller draft model. In this paper, we establish a surprising connection between speculative sampling and the concept of channel simulation from information theory, which aims at simulating a noisy channel using as few bits as possible. This connection allows us to provide an information-theoretic analysis of the speed up that can be achieved by speculative sampling. Leveraging this link, we derive an explicit relation between generation speed-up and the number of tokens k generated by the draft model for large k , which serves as an upper bound for all k . We also propose a novel speculative sampling method via exponential races called ERSS that matches state-of-the-art performance.

## introduction transformer-based large language models (llms) are at the forefront of the ai revolution, driving rapid advancements across numerous applications. however, as their adoption accelerates, the speed of text generation emerges as a critical bottleneck, alongside compute and memory constraints. standard llm generation involves calculating the conditional probability distribution using the target model p given a partial output x :@xmath0, x 2 , ..., x n ), which is the concatenation of the initial context and all previously generated tokens. the target model p outputs the conditional probability distribution p ( • | x :n ), from which a single token x n+1 is sampled and concatenated to x :n to form the updated sequence x :n+1 . these steps are repeated iteratively until a stopping condition is met, such as reaching a maximum sequence length or generating an end-of-sequence token. this auto-regressive token-by-token approach inherently leads to slow generation speeds. speculative decoding accelerates this process by generating multiple draft tokens for each target model evaluation using a more efficient draft model. these draft tokens are then verified by the target model, and some of them are chosen and appended to the input token sentence, while guaranteeing identical output quality; that is, the distribution of generated text remains the same as standard autoregressive decoding. in more detail, speculative sampling consists of the following three steps: speculative decoding achieves a 2-3× speed-up in text generation @xcite , the speed-up increasing with the alignment of the draft model q with the target model p . following common computational models, we quantify speed-up as the expected number of generated tokens per evaluation of the target model p . for a fair comparison of different speculative decoding strategies, we fix the number of drafted tokens k and analyze the speed-up achieved. while 'speculative decoding' describes the overall acceleration strategy for llms, the specific methods implementing these steps via probabilistic acceptance mechanisms are known as speculative sampling. channel simulation is a compression problem focused on efficiently generating samples from a target probability distribution p at a decoder @xcite . while only the encoder knows the target distribution p , both the encoder and decoder can access to samples from common reference distribution q. unlike standard source coding, the goal is not to communicate a specific sample, but rather to ensure that the decoder can generate any sample distributed according to p . common methods rely on generating a shared codebook (list) of samples at the encoder and decoder. based on the target distribution p , the encoder communicates to the decoder an index from this list, guaranteed to follow the target distribution p . a good introduction is provided by @xcite , and li (2024) compiles a comprehensive literature review. intuitively, both speculative sampling and channel simulation aim to generate samples from some target distribution p by generating samples from a draft/reference distribution q. in channel simulation, the goal is to minimize the entropy of the chosen index. in speculative sampling, the goal is to maximize the probability of one of the drafted tokens being accepted. although these objectives seem different, the most common speculative sampling method @xcite , which we refer to as greedy speculative sampling (gss), (explained in the next section), and a channel simulation method called greedy rejection sampling @xcite are essentially the same procedures! to explain this surprising connection, we develop a theory that explicitly links the entropy of the accepted token with the speed-up of speculative sampling by considering the generation of multiple tokens. additionally, we propose a novel speculative sampling strategy, exponential race speculative sampling (erss), based on another channel simulation method called poisson functional representation @xcite . we demonstrate that its performance matches that of gss. our contributions in this paper are: throughout this paper, ω denotes the finite set of tokens, where each token is represented by an integer corresponding to a textual element. sets are enclosed in braces {}, sequences in parentheses ( ). concatenation of sequences a and b is written as a||b. when clear from context, we write a||e as a shorthand for a||(e), where (e) is a sequence with a single element e. ## speculative sampling during inference, llms predict the distributions of tokens conditioned on the preceding sequence. we can consider an llm as a black box that takes a sequence of tokens x :n as input, and outputs a set of distributions {p ( • | x :i )} @xmath1. in standard auto-regressive decoding, we typically utilize only the last distribution p ( • | x :n ) to sample the next token. in contrast, speculative decoding leverages algorithm 1 simple gss (@xmath2 a straightforward extension to the simple drafting strategy is to speculate a sequence of k tokens instead of just one. during the drafting stage, the draft model q is run autoregressively k times to generate a draft sequence of length k, (x (1) , x(1,1) , . . . , x(1) k @xmath3. the notation xj describes that the token at position @xmath4, . . . , j m-1 , j m ) comes after token xj ′ where j ′ = (j 1 , . . . , j m-1 ). this choice of notation is made to ensure consistency across all strategies introduced later. the evaluation step then processes the entire drafted sequence x :n ||(x (1) , x(1,1) , . . . , x(1) k @xmath5. the drafted tokens are verified and accepted sequentially, starting from the first, until a rejection occurs, or all k drafted tokens are accepted. consequently, the sequence strategy can generate from 1 to k + 1 tokens per evaluation of the target model p . an alternative to sequence drafting is batch drafting @xcite , where multiple candidate continuations are explored. specifically, q is evaluated once to sample multiple draft token alternatives, which are then evaluated in parallel by the target model p . while naively computationally intensive, @xcite showed that it can be done at the same computational cost as sequence drafting. foot_0 this yields distributions p ( • | x :n ||x (i) ) @xmath6. the verification stage then sequentially considers each drafted token x(i) until the first acceptance. this strategy generates two tokens if any drafted token is accepted, or one if all get rejected. to avoid redundancy, candidates should be sampled without replacement @xcite . to recall the connection highlighted in section 1, the acceptance criterion in greedy rejection sampling @xcite and batch gss when @xmath7. the tree drafting strategy, which subsumes simple, sequence, and batch methods, operates on any ordered tree topology. each vertex represents a token drawn from q. each path from the root to a leaf constitutes a possible sequence drafted autoregressively. the general tree gss algorithm, detailed in appendix a and algorithm 4, recursively applies batch selection at each vertex of the tree to generate the accepted sequence. all the strategies are illustrated in figure 1 . the work most closely resembling ours is @xcite , which investigates optimal tree topologies for gss under specific assumptions. our analysis is more general and encompasses their problem setting as a special case. moreover, by establishing connections to information theoryspecifically channel simulation (li, 2024) and tunstall coding @xcite ) -we derive an upper bound on the expected speed-up, which is asymptotically exact as number of drafted tokens k increases. our algorithm for generating sampling trees achieves o(k log k) complexity, outperforming the o(k 2 |ω|) complexity of @xcite . additionally, we consider token-by-token generation, unlike global generation @xcite , which operates on trees natively. an alternative approach to speculative decoding leverages optimal transport (ot) principles @xcite @xcite . these methods frame the selection of accepted tokens as a linear programming problem. however, the dimensionality of this problem scales aa |ω| k , rendering it computationally intractable even for small values of k. consequently, approximations are employed to make ot-based methods practical. while these methods aim to maximize the immediate probability of accepting any single drafted token, this prioritization increases the entropy of accepted tokens. as we show in section 4, entropy of the acceptance distribution is inversely proportional to long-term speed-up (as k grows). in contrast, through the established connection to channel simulation, we show that both gss and our proposed erss implicitly aim at minimizing this acceptance entropy, thereby maximizing long-term generation speed-up. ## exponential races this section introduces the concept of exponential races. we first explain their use as a method for sampling from a distribution p . we then show their use for speculative sampling, leveraging a distribution q to 'predict' a sample from p . an exponential race @xcite is a poisson process with time-ordered points, each corresponding to a sample from a distribution p . we are interested in the winner (first point) of this race. in the discrete case, we can simulate relevant points, i.e., potential winners, by associating each element i of the sample space ω with an exponential random variable e i ∼ exp(1), @xmath8. the winner of the race, i * = arg min @xmath9, is distributed according to i * ∼ p . this is also known as the gumbel-max trick @xcite , where arrival times are obtained via the monotonic transformation -log(•). thus, given a distribution p , exponential races allow us to sample from it using independent exponential random variables. for completeness, the proof of correctness of this result is provided in appendix b. let p and q be distributions with the same support ω. using the same realization of e for exponential races yields i * @xmath10, both following their respective distributions. if p and q are similar (i.e., the ratio p i q i is close to 1), then it is likely that i * @xmath11. an example is illustrated in figure 2 , where q and p are distributions generated by llms. the simplest version of erss is presented in algorithm 2, where only a single token is generated from the draft model in the drafting stage: a winner of an exponential race under q( • | x :n ). as in gss, the target model p is then evaluated on the sequence x :n ||x. during verification, if the winner of the exponential race under p ( • | x :n ) is also x, one more token is drafted. it is straightforward to extend this approach to a sequence strategy, by generating an exponential race for every node in the draft sequence. for the batch strategy with k alternatives, the drafted tokens are selected as the first k arrivals of the exponential race under the draft model q, as those are the most likely winners of the race under law p . like the traditional gss, the batch case can be generalized to the full tree by continuing different trajectories. the general procedure for this tree approach is presented in appendix a as algorithm 5. we now highlight key distinctions between erss and gss. first, during verification, the fundamental operation of erss is finding the minimum of a vector, which is efficiently parallelizable. this contrasts with gss, which involves context: the cat sat on the e i q i i floor mat chair table sofa floor mat chair table sofa i e i p i an entirely sequential chain of accept/reject steps (see figure 1e ). second, erss verification can be implemented entirely in the unnormalized logit space. gss, however, requires the calculation of residual probability distributions and associated normalization constants at each step along its verification path. simplified computations due to its logit-only nature make erss applicable to continuous data, including images, video or audio, for use in variational autoencoders (kingma and welling, 2014) and diffusion models @xcite . third, the parallel minimum-finding mechanism lends erss greater conceptual simplicity compared to the stateful sequential logic of gss, potentially easing implementation and debugging. finally, regarding theoretical guarantees, section 4 presents analyses showing how erss's speed-up bounds (equation 8 ) compare favorably to those derived for gss (equation 7 ), implying tighter worstcase performance bounds for erss. ## markov chain simulation this section focuses on determining the optimal drafting strategy for speculative decoding. we derive this strategy by modeling the underlying distributions as markov chains. we then establish a connection to tunstall codes @xcite , which facilitates rigorous performance upper bounds for all speculative sampling methods and exact asymptotics for optimal speculative sampling. the central task is to identify the optimal drafting strategy, represented by a drafting tree τ . in such a tree, each vertex, except the root, is associated with a drafted token. specifically, for a vertex indexed by @xmath12, . . . , j l ), the associated token xj is sampled from q( • | x :n ||(x j :1 , . . . , xj :l-1 )), where j :@xmath13, . . . , j m ). let r( j | x :n ) denote the probability that vertex j, and thus the drafted token xj , is accepted during verification, given context x :n and a chosen speculative sampling algorithm. r( j | x :n ) is defined as the probability of acceptance, marginalized over all possible draft token sequences from q: for the rest of the section, we assume the target and draft model distributions, p and q, be m-th order markov sources. then, the acceptance probability r( j | x :n ) is an m-th order markov chain. the expected number of accepted tokens is the sum of acceptance probabilities for each vertex in the drafting tree. therefore, the optimal drafting tree τ * with k + 1 vertices (for speculating k tokens) is the solution to: crucially, for any vertex j and its descendant j ′ in a drafting tree, the acceptance probability is decreasing: r( j this is because j ′ can only be accepted if its ancestor j is accepted. consequently, the top k vertices with the highest acceptance probabilities form a valid tree. any node j in this top-k list will have all its ancestors j :m (where @xmath14. this observation leads to a greedy algorithm (algorithm 3) for constructing an optimal tree with k + 1 vertices. algorithm 3 iteratively builds the tree by greedily adding the next most likely token to be accepted. this is achieved efficiently using a priority queue to maintain candidate vertices, ordered by their acceptance probabilities. the algorithm's computational complexity is o(k log k) because the loop iterates algorithm 3 optimal tree construction 1: input: # drafted tokens k, partial output x :n 2: tree ← {()} ▷ initialize tree with root only 3: c ← priorityqueue( ) 4: add(c, (r((0) | x :n ), (0))) 5: for all @xmath15, . . . , k) do 6: _, j ← popmax(c) 7: our analysis relies on the acceptance probability function r, which is typically unknown in practice. if we approximate r using an empirical acceptance distribution -effectively treating it as a 0-th order markov source -the resulting algorithm becomes equivalent to that proposed by @xcite . we employ this empirical approximation in our experiments. however, as section 5 will demonstrate, this 0-th order approximation proves to be inaccurate, leading to an underestimation of the speed-up we can achieve with speculative decoding. developing more refined approximations of r presents a promising avenue for future investigation. in another connection to information theory, the process of generating optimal drafting trees (algorithm 3) closely resembles the construction of tunstall codes @xcite . due to their similarity, quantities in speculative sampling, such as the number of drafted tokens, speed-up, and entropy of the acceptance distribution, have direct counterparts in tunstall codes: the number of expanded nodes in the tunstall tree, the expected length of the consumed source symbols, and the source entropy, respectively. this connection allows us to express the expected speed-up of speculative sampling (both gss and erss) in terms of these quantities. tunstall code is a type of variable-to-fixed length source code used in the compression of discrete sources. in source coding, the goal is to represent sequences of symbols from a source alphabet (like tokens in our case) using the fewest number of bits (in the case of a binary alphabet). variable-to-fixed length codes, such as tunstall codes, achieve this by mapping variable-length sequences of source symbols to fixed-length codewords from the code alphabet. specifically, a tunstall code takes a variable-length prefix of the source symbol sequence and encodes it into a fixed-length output sequence. this encoding step is repeated until the entire source sequence is processed. tunstall codes are known to be optimal in the sense that for sufficiently long codewords, the average number of output alphabet symbols per source symbol approaches the entropy of the source -this also holds true for sources with memory @xcite . let the entropy of the acceptance probability distribution be defined as: definition 4.1. let s : p ×q×ω * ×@xmath16, draft model q, context x :n , and drafting strategy τ , and outputs a sequence of generated tokens. the expected number of tokens generated by s under strategy τ is defined as theorem 4.2. (upper bound tunstall) for speculative sampling employing the optimal drafting strategy τ * , the expected number of generated tokens is bounded as follows: the proof of theorem 4.2 is provided in appendix d. since this bound applies to the optimal strategy τ * , it holds for any strategy τ . as the theorem suggests, for a large number of drafted tokens k, the generation speed-up g τ * is fundamentally governed by the entropy of the acceptance distribution, h [r]. channel simulation algorithms, such as greedy rejection sampling and poisson functional representation, are explicitly designed to minimize this entropy. consequently, these methods, which underpin both gss and erss, inherently aim to maximize generation speed-up by reducing h @xcite . from a channel simulation perspective, the for any ϵ > 0, and for poisson functional representation (li and el gamal, 2018) (relevant to erss): these bounds show the connection of speedup in speculative sampling to kl divergence d kl [p ∥q] between the models. while asymptotically accurate for large k, the bound from theorem 4.2 is dominated by the full token alphabet size |ω| when the number of drafted tokens k is small. for such values of k, we observe that the optimal drafting tree typically explores only a limited number of drafting positions. indeed, multiple acceptance distributions can yield identical optimal drafting trees and the same acceptance probabilities for all nodes within those trees. thus, evaluating the bound from theorem 4.2 for any such distribution provides a valid upper bound for all of them. to simplify the notation and exposition, we focus on the 0-th order acceptance distribution r, noting that our findings can be generalized to higher-order markov sources. for a given number of drafted tokens k, let d be the maximum index explored in the optimal drafting tree τ * k with k + 1 vertices, i.e., @xmath17. we define an equivalent acceptance distribution r such that for indices @xmath18, r(i) = r(i), and for @xmath19, r(i) ≤ min i ′ ≤d r(i ′ ). in short, r matches r on tree nodes for the first d indices and is upperbounded by min i ′ ≤d r(i ′ ) thereafter. our objective is to minimize the upper bound on the expected number of accepted tokens. upon inspection of theorem 4.2, we see that a tighter bound is achieved by maximizing the entropy of r and minimizing the alphabet size. minimizing the bound requires strategically distributing the remaining probability mass, p @xmath20, associated with indices beyond d. for any choice lemma 4.3. (upper bound m) for a fixed number of drafted tokens k and initial index range d, consider a 0-th order acceptance probability distribution r. let p @xmath21. define r res such that r res (i) = r(i) for @xmath22, and r res (d + 1) = p res . then, for any integer @xmath23, the expected number of generated tokens is upper-bounded by: note: r res defined in above lemma is not r. ## experiments to validate our theoretical analysis, we conducted numerical experiments comparing gss and erss across different drafting strategies. we performed open-ended text generation up to 200 tokens, accumulating 100k generated tokens per strategy. the experiments were performed on 8 nvidia rtx a6000 gpus, with experiments running for 140 hours of wall-clock time. we estimated the acceptance probability function r using empirical acceptance probabilities for different indices (from sequence and batch ), approximating it as a 0-th order markov chain. based on this estimated r, we computed the optimal drafting tree τ * for both gss and erss using algorithm 3. figure 3 illustrates the expected number of accepted tokens as a function of drafted tokens for the target model llama-3.1 70b and the draft model llama-3.2 1b @xcite , used under the meta llama 3.1 and 3.2 licenses. figure 4 shows the marginal change in acceptance probability with each additional drafted token. we compare four drafting strategies: batch, sequence, optimal tree (τ * ), and the specinfer tree @xcite , which drafts 3-sequences with the first two drafted tokens being common. for each strategy and speculative sampling method (gss, erss), we present both theoretical and empirical performance. the theoretical plots show the expected number of generated tokens based on the estimated r, while the empirical plots display the actual value observed in our experiments. figure 3 reveals that for batch drafting with @xmath24, gss outperforms erss-it has a higher chance of accepting the first proposed sample. this can be explained by the 'greediness' of gss, which maximizes the acceptance probability of each sample sequentially, a strategy that is suboptimal when considering multiple samples. this also explains the performance gap between strategies for small values of drafted samples k. further discussion regarding the probability of acceptance of the first sample for gss and erss is presented in appendix c. however, with @xmath25, both methods exhibit comparable performance. for @xmath26, erss achieves a higher expected number of accepted tokens, and hence speed-up. the acceptance probabilities for each index ((0), (1), (2), . . . ) are shown in figure 4 . for sequence drafting, gss demonstrates superior performance compared to erss, since this strategy comprises a series of batch drafting steps with @xmath27, considering only a single possible continuation at each step. intriguingly, our empirical results also challenge the 0-th order acceptance assumption @xcite . the empirical performance plateaus at a higher level than predicted by the measured r, indicating higher-order markov dependencies. this can be explained by regions of language where the target model p and draft model q exhibit greater alignment, leading to extended sequences of accepted tokens. the specinfer tree achieves slightly improved performance over sequence drafting but still plateaus. as predicted, the optimal τ * tree strategy yields the best performance for both gss and erss, exhibiting a logarithmic relationship between drafted tokens k and expected generated tokens (figure 3 ). at low k, the optimal tree τ * and sequence strategies show similar performance, with gss slightly outperforming erss. however, as k increases, the performance gap diminishes, and they converge. figure 3 presents the minimal upper bound derived from lemma 4.3 by optimizing the parameter m for each number of drafted tokens k. we also depict the upper bound from lemma 4.3, calculated with a fixed m for @xmath28. while the marginal changes in the minimal upper bound are omitted for visual claritydue to their step-like transitions-, we observe that the marginal changes of the fixed-m upper bound closely follow the trends of the optimal τ * tree . ## future directions our channel simulation perspective on speculative sampling suggests several avenues for future research. first, beyond exact sampling which guarantees output identical to the target model p , one could explore approximate speculative sampling. drawing from work on approximate channel simulation @xcite , which relates approximation quality to communi-cation cost (analogous to draft sample count), this could illuminate the trade-offs between speed-up and distribution shifts, potentially leading to new algorithms. second, further application of the poisson matching lemma (li and anantharam, 2021), utilized for our lemma c.1 upper bound, could yield lower bounds on acceptance probability or index entropy for erss, complementing our analysis with worst-case performance guarantees for speed-up. third, the channel simulation framework links speed-up to the accepted index entropy. this motivates comparing different draft model q optimization strategies: directly minimizing this index entropy versus minimizing d kl (p ||q), or using reinforcement learning objectives targeting empirical speed-up. finally, while prior work has connected speculative decoding to ot, our channel simulation framework suggests investigating minimal entropy couplings as a more direct theoretical foundation for maximizing long-term speed-up, distinct from standard ot objectives. ## conclusion this work establishes a connection between speculative decoding -a technique for accelerating autoregressive llm generation -and channel simulation. this connection enabled us to propose erss, a novel speculative sampling method. by linking the optimal drafting strategy to tunstall codes, we derived a theoretical upper bound and the asymptotic relationship between the number of speculated tokens and the expected speed-up, for gss and erss. these findings offer a deeper understanding of, and potential improvements to, the efficiency of speculative decoding.