Faster Cascades via Speculative Decoding

Harikrishna Narasimhan${}^{†}$, Wittawat Jitkrittum${}^{\ast }$, Ankit Singh Rawat${}^{\ast }$
Seungyeon Kim${}^{\ast }$, Neha Gupta${}^{†}$, Aditya Krishna Menon${}^{\ast }$, Sanjiv Kumar${}^{\ast }$

${}^{†}$Google Research, Mountain View ${}^{\ast }$Google Research, New York

Corresponding author: [email protected]

Cascades and speculative decoding are two common approaches to improving language models' inference efficiency. Both approaches involve interleaving models of different sizes, but via fundamentally distinct mechanisms: cascades employ a deferral rule that invokes the larger model only for "hard" inputs, while speculative decoding uses speculative execution to primarily invoke the larger model in parallel verification mode. These mechanisms offer different benefits: empirically, cascades offer better cost-quality trade-offs, often even outperforming the large model, while theoretically, speculative decoding offers a guarantee of quality-neutrality. In this paper, we leverage the best of both these approaches by designing new speculative cascading techniques that implement their deferral rule through speculative execution. We characterize the optimal deferral rule for our speculative cascades, and employ a plug-in approximation to the optimal rule. Experiments with Gemma and T5 models on a range of language benchmarks show that our approach yields better cost-quality trade-offs than cascading and speculative decoding baselines.

Large language models (LLMs) have yielded significant advances in quality on a range of natural language processing tasks ([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]), at the cost of an increase in inference latency. This has sparked a growing body of literature on reducing LMs' inference costs without (overly) compromising on quality ([12, 13, 14, 15, 16, 17, 18]). One such line of work involves constructing a family of models of various sizes (e.g., a small and large model), and suitably orchestrating amongst them to make a prediction. Two canonical instantiations of this strategy are model cascading ([19, 20, 21, 22, 23, 24, 25, 26]) and speculative decoding ([27, 16, 15, 18, 28, 29]).

While similar in spirit, cascades and speculative decoding are fundamentally different in details. Cascades employ a deferral rule to identify "hard" inputs, and only invoke larger models on such inputs. For example, in a two-model cascade, one first invokes the smaller model, and uses its associated probability of the generated output to decide whether to defer to the larger model. By contrast, speculative decoding uses a small model to draft a block of tokens via standard auto-regressive decoding, which are then verified in parallel by a large model. One then accepts all drafted tokens until the first "implausible" one, which is rolled back based on the larger LM's prediction.

Owing to their different mechanisms, both methods have complementary strengths. Cascades seek to output distributions that have the best quality for a given cost budget, and potentially provide * better cost-quality trade-offs*, sometimes even yielding better accuracies than the individual models they are constructed with ([30, 31]) (§ 3). By contrast, speculative decoding is theoretically guaranteed to match the output distribution (or a close approximation thereof ([32])), and are practically observed to provide impressive speed-ups ([27, 16, 15, 18]). Given the complementary nature of these two approaches, a natural question arises: can we leverage the best of both techniques?

In this paper, we do so by designing new techniques for two-model cascades that implement their deferral rule in a speculative manner: we have the smaller model generate drafts auto-regressively, and the larger model execute in parallel on the drafts to decide whether or not to defer on them. We show that this speculative cascading approach yields better cost-quality trade-offs than both standard cascades and speculative decoding. In detail, we make the following contributions:

Lossy speculative sampling ([32]) is a special case of this recipe for a particular target distribution (§ 4.1). 2. We show how common cascading rules, such as Chow's rule ([33]) and confidence-difference thresholding ([30]), can be implemented speculatively by plugging in their target distribution into our framework. We refer to these as speculative cascades (§ 4.2). 3. We characterize the theoretically optimal deferral rule for a speculative cascade, and design a speculative cascading technique that implements a plug-in estimate to the optimal rule (§ 4.3, Lemma 6, Table 1). We also present token-specific variants of our deferral rules (§ 4.4). 4. Through experiments with Gemma ([34]) and T5 models ([35]) on a range of benchmark language tasks including summarization, translation, reasoning, coding and QA, we show that speculative cascades are able to provide better cost-quality trade-offs than their sequential cascade and speculative decoding counterparts (§ 6).

Let $\mathcal{V}$ denote a finite vocabulary of tokens, with ${\mathcal{V}}^{\ast }$ denoting the set of all sequences generated by this vocabulary. Let ${\Delta }_{\mathcal{V}}$ denote the set of all probability distributions over tokens in $\mathcal{V}$. Given an arbitrary length sequence $x=x_1{x}_2\dots x_L\in {\mathcal{V}}^{\ast }$ and index $i\le L$, denote by $x_{<i}=x_1{x}_2\dots x_{i-1}$. A language model (LM) is a probability distribution over ${\mathcal{V}}^{\ast }$. Let $\mathbb{P}$ denote the ground-truth probability distribution over ${\mathcal{V}}^{\ast }$. This could be, for example, a distribution over prompt-response pairs that the LM may encounter during deployment, or a distribution of sequences used to pre-train the LM. We will measure the quality of an LM based on how closely it mimics $\mathbb{P}$.

Suppose we are provided two LMs $q$ and $p$, where $p$ is the larger (more expensive) model. Our goal is to design an inference strategy that selectively invokes $q$ and $p$ to trade-off between quality and latency (which may be approximated by the fraction of times that $p$ is invoked). We will denote by $q(x_t|x_{<t})$ the probability $q$ associates to token $x_t\in \mathcal{V}$ given prefix $x_{<t}\in {\mathcal{V}}^{t-1}$, and by $p(x_t|x_{<t})$ the same distribution from model $p$. Whenever it is clear from context, we will hide the conditioning on prefix $x_{<t}$, and use the shorthand $p_t(\cdot )$ for $p(\cdot |x_{<t})$ and $q_t(\cdot )$ for $q(\cdot |x_{<t})$.

Cascades are an effective strategy to trade-off cost and quality by having the smaller model $q$ handle the "easy" samples, and the larger model $p$ handle the "hard" ones ([25, 36]). A common cascading approach is confidence thresholding or Chow's rule ([33, 30]), where we first run $q$ on the input, and defer to $p$ when $q$ 's confidence for its generated response is sufficiently low. This strategy is typically implemented at the sequence-level, where for a given prefix $x_{<m}$ we invoke $q$ to generate a complete response $x_m\dots x_{m+n}$. We evaluate $q$ 's predicted probability for the response, and check whether it falls below a threshold $\alpha \in [0,1]$:

If the above holds, we defer to $p$ to generate a new response; otherwise, we retain $q$ 's response. One may tune the threshold to achieve a desired cost-quality trade-off. The literature also offers variants of Chow's rule that use a more nuanced aggregation of per-token uncertainties ([25]).

Speculative decoding is an alternate strategy that applies token-level interleaving between $q$ and $p$, resulting in provably matching the larger model quality at a reduced inference cost ([27, 15]). Given a prefix $x_{<t}$, we draft $\gamma$ draft tokens $x_t,\dots ,x_{t+\gamma -1}$ via auto-regressive sampling from $q$, and verify if these tokens can be accepted by running $p$ in parallel on the $\gamma$ prefixes $x_{<t},\dots ,x_{<t+\gamma -1}$. We then rollback to the first rejected token $t+j^{\ast }$ (where $j^{\ast }\in \{0,1,\dots ,\gamma -1\}$), replace $x_{t+j^{\ast }}$ with a new token, and repeat the process with prefix $x_{<t+j^{\ast }+1}$.

During the verification stage, a draft token $x_{t+j}$ generated by $q$ is accepted with probability $\min \left(1,\frac{p_{t+j}(x_{t+j})}{q_{t+j}(x_{t+j})}\right)$ and rejected otherwise, recalling the shorthand $q_{t+j}(\cdot )=q(\cdot |x_{<t+j})$ and $p_{t+j}(\cdot )=p(\cdot |x_{<t+j})$. A rejected token is then replaced by a new token sampled from a modified distribution $\mathrm{norm}\left(\max \left\{0,p_{t+j}(\cdot )-q_{t+j}(\cdot )\right\}\right),$ where $\text{norm}(\cdot )$ denotes normalization to sum to 1. This sampling process is provably equivalent to sampling $\gamma$ tokens auto-regressively from $p$ for prefix $x_{<t}$ ([15]). We summarize this speculative sampling procedure in Algorithm 1. Each invocation of this algorithm generates at most $\gamma +1$ next tokens (and at least one) for a given prefix $x_{<t}$. One may run this algorithm multiple times to generate a complete output sequence.

In practice, one may employ a lossy variant ([32]) of the above sampling that allows some deviation from verifier's distribution $p$. In this case, a draft token $x_{t+j}$ is accepted with probability $\min \left(1,\frac{p_{t+j}(x_{t+j})}{(1-\alpha )\cdot q_{t+j}(x_{t+j})}\right)$, where $\alpha \in [0,1)$ is a strictness parameter, with higher values indicating greater deviation from $p$. A rejected token may then be replaced by a token sampled from the residual distribution $\mathrm{norm}\left(\max \left\{0,\frac{1}{\beta }\cdot p_{t+j}(\cdot )-q_{t+j}(\cdot )\right\}\right),$ where $\beta \ge 1-\alpha$ is a parameter that depends on $\alpha$, $q$ and $p$. A common heuristic is to simply set $\beta =1$ ([37]).

Both cascades and speculative decoding interleave models of different sizes to reduce inference cost, but fundamentally differ in the mechanisms they use. As a step towards comparing the strengths and weaknesses of these approaches, we first describe how one may design a token-level cascade.

It is straightforward to extend the sequence-level Chow's rule from § 2 to form a token-level cascade between $q$ and $p$. For a prefix $x_{<t}$, we first compute the smaller model's distribution $q(\cdot |x_{<t})$, and check whether ${\max }_{v\in \mathcal{V}}q(v|x_{<t})$ is below a pre-chosen threshold. if so, we evaluate $p(\cdot |x_{<t})$, and sample $x_t\sim p(\cdot |x_{<t})$; otherwise, we sample $x_t\sim q(\cdot |x_{<t})$.

More generally, we may design a token-level deferral rule $r:{\mathcal{V}}^{t-1}\to \{0,1\}$ that takes the prefix $x_{<t}$ as input and outputs a binary decision, with $r(x_{<t})=1$ indicating that we defer to $p$ (i.e., draw a sample from $p$ rather than $q$). For example, token-level Chow's rule can be written as:

where $\alpha$ is a threshold parameter; the higher the value, the lower is the frequency of deferral to $p$. One may also use other confidence measures than the maximum probability, such as the entropy of the small model's probability distribution. We elaborate in § B that the choice of confidence measure would depend on the evaluation metric of interest; Equation 2 is typically prescribed when the cascade's quality is evaluated in terms of its accuracy against the ground-truth distribution on individual tokens, whereas entropy is prescribed when the metric of interest is the cross-entropy loss.

While Chow's rule Equation 2 is easy to implement, it can be sub-optimal if the smaller model's max-token probability is not reflective of which of the two models are better equipped to predict the next token for a given prefix ([30]). Given this, it is natural to ask what the optimal deferral rule $r$ for a token-cascade looks like, and whether we can reasonably approximate this rule.

For this, we must first specify an objective to minimize at each step $t$. Following the prior cascade literature ([30, 25]), a reasonable objective to minimize is the expected loss from the deferral rule against the ground-truth distribution $\mathbb{P}$, with an added cost for deferring to the larger model. We state this below for a fixed prefix $x_{<t}$, using as before the short-hand $q_t(\cdot )$ for $q(\cdot |x_{<t})$ and $p_t(\cdot )$ for $p(\cdot |x_{<t})$:

for a cost penalty $\alpha \ge 0$ and loss function $\ell :\mathcal{V}\times {\Delta }_{\mathcal{V}}\to {\mathbb{R}}_{+}$. Common choices for $\ell$ include the 0-1 loss ${\ell }_{\text{0-1}}(v,q_t)=1\left(v\ne {\mathrm{argmax}}_{v^{\prime }}{q}_t(v^{\prime })\right)$ and the log loss ${\ell }_{\log }(v,q_t)=-\log \left(q_t(v)\right).$

The minimizer of Equation 3 is of the form:

Intuitively, we compare the expected loss from $q$ with the expected cost of invoking $p$, and decide to defer when the latter is smaller. We note here that this optimization problem is set up for a fixed prefix $x_{<t}$. One may also consider the coupled optimization problem across all positions.

Plug-in estimator for Equation 4. The optimal rule in Equation 4 requires computing expectations over the ground-truth distribution $\mathbb{P}(\cdot |x_{>t})$, which is not available during inference time. A common approach in the cascades literature is to replace the expected losses with the models' confidence estimates ([30]). For example, when $\ell ={\ell }_{\text{0-1}}$, it may be reasonable to use $1-{\max }_v{q}_t(v)$ as an estimate of the expected 0-1 loss ${\mathbb{E}}_{x_t\sim \mathbb{P}(\cdot |x_{<t})}\left[{\ell }_{\text{0-1}}(x_t,q_t)\right]$ and $1-{\max }_v{p}_t(v)$ as an estimate of ${\mathbb{E}}_{x_t\sim \mathbb{P}(\cdot |x_{<t})}\left[{\ell }_{\text{0-1}}(x_t,q_t)\right]$. The extent to which these estimates are accurate depend on how well $q$ and $p$ are calibrated ([38]). The resulting plug-in estimator for (Equation 4) thresholds the difference of confidence estimates from both distributions:

Similarly, when $\ell ={\ell }_{\log }$, we may use the entropy $-{\sum }_v{q}_t(v)\cdot \log (q_t(v))$ from $q_t$ as an estimate of its expected log-loss, and similarly for $p_t$ (see Appendix C).

For efficiency reasons, ${\hat{r}}_{\mathtt{Diff}}$ cannot be directly used in a token-level cascade, as it needs the large model to be invoked at every step $t$.

However, it serves as an oracle that allows to analyze the head-room available to improve upon Chow's rule.

See also Remark 4.

Token-level cascades and speculative decoding differ in the distribution over tokens they seek to mimic. Speculative decoding seeks to mimic the large model's output distribution, and is usually used when one wants to match the quality of the large model. On the other hand, token-level cascades seek to output distributions that closely approximate the ground-truth label distribution and potentially offer * good cost-quality trade-offs*, sometimes yielding better quality than even the large model.

Cascades are useful when the draft model fares better than the verifier on some inputs, and one may want to retain the drafter's predictions even when it disagrees with the verifier. Even in cases where both the drafter and verifier fare poorly on some inputs (e.g., due to label noise), one may want to ignore the disagreement between the drafter and verifier to avoid triggering unnecessary roll-backs.

As a concrete example, we consider token-level cascades of T5 models ([35]) of two different sizes finetuned on a WMT EN $\to$ DE translation [39] and an extreme summarization (XSum) task ([40]). We construct these cascades using both (token-level) Chow's rule in Equation 2 and the oracle Diff rule in Equation 5, and also apply speculative decoding with the smaller (larger) model as the drafter (verifier). In Figure 1, we plot quality as a function of fraction of samples deferred to the large model (number of deferrals divided by number of generated tokens), as we vary the cost parameter $\alpha$. Note that with speculative decoding, each verification step verifies $\gamma$ tokens in parallel, but is counted as a single deferral to the large model. While speculative decoding matches the quality of the large model (right-most point), the oracle yields significantly better cost-qualty trade-offs. Even Chow's rule, which is sub-optimal for cascading ([30]), offers better trade-offs, and outperforms speculative decoding in a small region. As noted by [31], this may be attributed to the ensembling effect in a cascade.

However, as also evident from the plots, token-level cascades require a significantly larger number of deferrals to the large model to achieve the same quality. This is because token-level cascades are executed sequentially: whenever $q$ defers, we execute $p$ once to generate one next token for the prefix accumulated so far, and the control transfers back to $q$. In contrast, speculative decoding runs $p$ in scoring mode to verify $\gamma$ draft tokens from $q$ in parallel. Moreover, the stochastic verification algorithm in speculative decoding often results in fewer tokens from $q$ getting rejected compared to the deterministic deferral rules used in a cascade. These observations motivate a natural question: given their complementary strengths, how can we leverage the best of both these techniques?

In addressing the above question, we present our main contribution: a principled approach to combining the better trade-offs cascades offer with the faster execution of speculative decoding.

We begin by considering a generic version of speculative sampling that seeks to mimic a general target distribution derived from the drafter's and verifier's distributions. In the proposed sampling procedure outlined in Algorithm 4, we sample tokens auto-regressively as before from the drafter's distribution. During the verification step, however, we do not compare the drafter's token probabilities against the verifier's distribution. Instead, we use a user-specified target distribution $\pi =\mathbb{T}(q,p)\in {\Delta }_{\mathcal{V}}$ derived from the drafter's and verifier's distributions at position $t$, for some function $\mathbb{T}(\cdot ,\cdot )$ that is inexpensive to compute. We accept a draft token $x_t$ when $q(x_t)\le \pi (x_t)$ and reject it otherwise with probability $1-\frac{\pi (x_t)}{q(x_t)}$. Upon rejection, we re-sample from the residual distribution $\mathrm{norm}\left(\max \{0,\pi (\cdot )-q(\cdot )\}\right)$.

This general procedure not only encompasses standard speculative decoding ([15]) for $\mathbb{T}(q,p)=p$, but also includes lossy speculative decoding ([32]) as a special case:

Algorithm 4 reduces to the lossy speculative sampling procedure in ([32]) with parameters $\alpha$ and $\beta$ when $\mathbb{T}(q,p)(v)=\max \{\min \{q(v),\frac{p(v)}{1-\alpha }\},\frac{p(v)}{\beta }\}$.

Equipped with Algorithm 4, we now propose new cascading techniques that implement their deferral rule in a speculative manner. Recall from § 3.1 that a token-level cascade of two models $q$ and $p$ is defined by a deferral rule $r:{\mathcal{V}}^{t-1}\to \{0,1\}$. For a prefix $x_{<t}$, the next-token distribution at position $t$ modeled by this cascade can be written as:

In fact, for all the deferral rules described in § 2, the resulting distribution can be described by a target distribution function ${\mathbb{T}}_{\delta }$ of the form:

for some function $\delta :{\Delta }_{\mathcal{V}}\times {\Delta }_{\mathcal{V}}\to \{0,1\}$ that maps distributions $(q,p)$ to a binary decision. For example, for Chow, $\delta (q,p)=1({\max }_{v}q(v)<1-\alpha )$, and for Diff, $\delta (q,p)=1({\max }_{v}q(v)<{\max }_{v}p(v)-\alpha ).$ See Table 1 for a summary of target distributions for different deferral rules.

Our proposal is to then invoke the speculative sampling procedure in Algorithm 4 with ${\mathbb{T}}_{\delta }$ as the target distribution function. We outline this generic speculative cascading approach in Algorithm 5, and contrast it with the sequential execution of a deferral rule in Algorithm 2.

In a sequential cascade, the large model's distribution $p$ cannot be used at the time the deferral decision is made (see Remark 2), as this would defeat the purpose of the cascade. With a speculative cascade, however, we can employ rules like Diff that depend on both $q$ and $p$. This is because we run the large model $p$ in parallel on drafts generated by the small model $q$, allowing us to compute both $p(\cdot )$ and $q(\cdot )$ on every prefix.

So far we have considered deferral rules designed for sequential cascades. In what follows, we derive the optimal deferral rule $r$ for a speculative cascade, where we sample speculatively from a target distribution $\pi =(1-r(x_{<t}))\cdot q_t+r(x_{<t})\cdot p_t$ using $q_t$ as the drafter.

As with sequential cascades (§ 2), we begin by defining an objective to minimize. We seek a deferral rule $r:{\mathcal{V}}^{t-1}\to \{0,1\}$ that minimizes a loss against the ground-truth distribution, while limiting the inference cost to be within a budget. (Per above, this deferral rule implicitly defines a target distribution $\pi$.) The inference cost crucially depends on how frequently a draft token is rejected in the verification phase, triggering a rollback. To this end, we derive the probability that a token sampled from $q$ is rejected during verification, for a target distribution resulting from a deferral rule $r$.

For a given prefix $x_{<t}$, and target distribution $\pi =(1-r(x_{<t}))\cdot q_t+r(x_{<t})\cdot p_t$, the probability of a token drawn from draft distribution $q_t$ being rejected is equal to: $r(x_{<t})\cdot D_{\text{TV}}(p_t,q_t),$ where $D_{\text{TV}}(p,q)={\sum }_{v\in \mathcal{V}}\max \{0,p(v)-q(v)\}$ is the TV distance between $p$ and $q$.

Intuitively, whenever $r(x_{<t})=0$, $\pi (v)=q_t(v)$, and therefore there is no rejection or roll-back; when $r(x_{<t})=1$, the rejection rate equals $D_{\text{TV}}(p_t,q_t)$.

For a fixed prefix $x_{<t}$, we formulate the goal of finding a solution to:

for some budget $B>0$. Equivalently, one may minimize an unconstrained objective similar to Equation 3, for suitable cost parameter $\alpha >0$ (see § C.4):

Contrasting Equation 8 with the deferral risk in Equation 3 for a sequential cascade, a key difference is that the cost of deferring to the larger model is no longer a constant, but depends on the similarity between $q_t$ and $p_t$, as measured by the TV distance between them.

We next derive the optimal deferral rule for Equation 8, and construct a feasible estimator for it.

The minimizer of Equation 8 is of the form:

When $p_t$ and $q_t$ are similar, the rejection rate for $q_t$ is low, and hence the deferral decision will depend largely on which of the two models yields a lower expected loss. When $p_t$ and $q_t$ are very different, the optimal decision is to defer to $p_t$ only when it yields a substantially lower loss than $q_t$.

Plug-in estimator for Equation 9. The optimal rule requires estimating expectations with respect the ground-truth distribution $\mathbb{P}(\cdot |x_{<t}).$ We employ similar plug-in estimators as the ones used with sequential cascades (§ 3). When $\ell ={\ell }_{\text{0-1}}$, we replace the expected 0-1 loss with (one minus) the maximum probability from the model, giving us:

The efficacy of the plug-in estimator depends on how closely the individual models approximate the ground-truth distribution $\mathbb{P}(\cdot |x_{<t})$; this is formalized by the following regret bound:

Suppose $\ell ={\ell }_{\text{0-1}}$. Then for a fixed prefix $x_{<t}$:

One can now run the speculative cascading procedure in Algorithm 5 using Equation 10 as the deferral rule; the corresponding $\delta (\cdot )$ is listed in Table 1. See § C.2 for a similar derivation for $\ell ={\ell }_{\log }$.

The plug-in deferral rules in (Equation 5) and (Equation 10) decide between the drafter's distribution $q_t(\cdot )$ and the verifier's distribution $p_t(\cdot )$ by comparing their maximum token probabilities. A downside to this approach is that the draft token $x_t\sim q_t(\cdot )$ may not maximize $q_t(\cdot )$. Thus, even when $x_t$ is of poor quality, we may end up accepting it because $q_t$ happens to be more peaked than $p_t$.

To alleviate this problem, we propose the use of token-specific deferral rules $r:{\mathcal{V}}^{t-1}\times \mathcal{V}\to \{0,1\}$ that use both the prefix $x_{<t}$ and a candidate token $v$ to provide a binary decision $r(x_{<t},v)\in \{0,1\}$, with 0 indicating that the token is of acceptable quality. We may then construct a target distribution of the following form:

where $\eta ={\sum }_{v^{\prime }\in \mathcal{V}}r(x_{<t},v^{\prime })\cdot q_t(v^{\prime })$ is a normalizing term chosen to ensure that ${\sum }_v{\pi }_{\mathtt{Token}}(v)=1$. This target distribution closely mimics $q_t(\cdot )$ on tokens that the deferral rule $r$ deems to be of acceptable quality, and defers to $p_t(\cdot )$ otherwise. One can modify the generic speculative sampling algorithm in Algorithm 4 to use ${\pi }_{\mathtt{Token}}$ as the target distribution, as shown in Algorithm 6 in § D.

To design $r$, we propose a heuristic variant of the Diff rule in Equation 4 that compares the expected 0-1 loss from the candidate token $v$ with the expected 0-1 loss from distribution $p_t$ (in § D, we discuss deriving a similar variant of the OPT rule in Equation 9):

for a cost parameter $\alpha$. The following are some simple plug-in approximations to Equation 12:

where we approximate $\mathbb{P}(v|x_{<t})$ with either $q_t(v)$ or $p_t(v)$. Equation 15 is a multiplicative plug-in approximation that has similarities to the rejection criterion used by [15] for lossy speculative greedy decoding, and results in an intuitive target distribution:

where ${\text{Top}}_{\alpha }=\{v\in \mathcal{V}:p_t(v)\ge {\max }_{v^{\prime }}{p}_t(v^{\prime })\cdot (1-\alpha )\}$ is the set of top ranked tokens by $p_t(\cdot )$. For these top-ranked tokens, ${\pi }_{TokenV3}$ approximates $q_t(\cdot )$; for the rest, it is a re-scaled version of $p_t(\cdot )$.

There has been a stream of work on improving the draft generation process in speculative decoding; these include having the drafter and verifier share the same backbone ([27, 41, 42, 43, 44, 45, 46, 47]), using multiple small draft models [48, 49], using tree-structured draft batches ([50, 51]), distilling the drafter with the verifier ([37]), and leveraging multiple sampled draft candidates [18].

The work that is most closely related to our specific proposal is the Big Little Decoder (BiLD)~([31]), which can be seen as another lossy variant of speculative decoding ([15, 32, 37]). BiLD has two phases: a fallback phase, during which the drafter $q$ is run auto-regressively until its maximum predicted probability is sufficiently low; and a rollback phase, during which the verifier $p$ is run in parallel on the prefixes generated by $q$ and rolls back to the point where $D(q,p)>\alpha$, for a metric $D$ that measures discrepancy and threshold $\alpha$. The fallback phase implements Chow's deferral rule in (Equation 2), and allows for the draft window size to vary dynamically based on an estimate of how likely the draft tokens will be accepted; the rollback phase can be seen as a deterministic variant of the rejection sampling algorithm of [15].

An advantage of BiLD over the rejection sampling algorithm in ([15]) is the use of Chow's rule to vary the draft window size. However, the final target distribution it seeks to mimic, ${\mathbb{T}}_{\text{BiLD}}(q,p)(v)=1(D(q,p)\le \alpha )\cdot q(v)+1(D(q,p)>\alpha )\cdot p(v)$, is an approximation to $p$; specifically, the target distribution $\pi ={\mathbb{T}}_{\text{BiLD}}(q,p)$ is chosen to satisfy $D(\pi ,p)\le \alpha$. Hence, in cases where $q$ deviates substantially from $p$, BiLD would choose $p$ as the target distribution, even when $q$ offers better quality on a prefix (where quality can be measured using a suitable loss function). In contrast, our proposed approach in § 4 uses speculative decoding to approximate target distributions that seek to optimally cascade between $q$ and $p$. In our experiments, we compare the efficacy of using ${\mathbb{T}}_{\text{BiLD}}$ as the target distribution with the target distributions we propose in this paper (see Table 1).

We compare our speculative cascading techniques with both sequential cascades and standard speculative decoding on a range of language benchmarks, including translation, reasoning, coding, QA, etc. We evaluate speculative cascades constructed from both the T5 v1.1 family of encoder-decoder models ([2]), and Gemma v2 decoder-only models ([34]). We construct the cascades with four deferral rules: (i) Chow in (Equation 2), (ii) Diff in (Equation 5), (iii) OPT in (Equation 10), and (iv) the Token-specific rule in (Equation 15) (we present results for the V1 and V2 variants in § E.7).

Cascades versus SpecDecode evaluation. Our evaluation protocol is markedly different from the standard evaluation of speculative decoding algorithms, where the goal is to speed up inference with a large model while preserving its output distribution. In contrast, our focus is on trading-off quality for lower inference costs by cascading two models of different sizes. We also do not claim to develop a new state-of-the-art method for fast LM inference. Furthermore, the speculative cascades we design build on the original speculative decoding algorithm [15]. While one could potentially also adapt our proposal to other recent variants of speculative decoding ([42, 52]), these involve a wholly orthogonal suite of techniques to what we propose (such as architectural changes, allowing for multiple drafts, distillation, and so on; see § 5).

Baselines. The cascading and speculative decoding methods we compare to include:

Fine-tuned T5 cascades. Our experiments on T5 models are based on the setup in [37]; see § E.1 for details. We use T5-small (77M) as the small model, and either T5-large (800M) or T5-XL (3B) as the large model. In each case, we supervised fine-tune these models on three tasks: WMT EN $\to$ DE translation ([54]), CNN/DM summarization ([55]), and XSum abstractive summarization ([40]). We use temperatures $T=0,0.1,0.5,1.0$, and block sizes $\gamma =3,5,7$ (full results in § E). Following the protocol in [15, 37], to measure latency, we evaluate the wall-clock decoding time with batch size 1.

In Figure 2, we present plots of quality vs. latency for the different methods. In each case, we vary the lenience parameter $\alpha$, and plot either the BLEU or ROUGE-2 metric as a function of the relative latency to the larger model. For brevity, we include the three main baselines; in § E.5–Appendix E.6, we compare to SpecDecode [Lossy${}^{\star }$] ([32]) and the original BiLD algorithm [31]. Methods that use speculative execution are considerably faster than sequential token-level cascades (TokenCascade [Chow]), although sequential cascades do have an advantage in the low-latency regimes. This is because unlike speculative approaches, which always call the large model after every $\gamma$ steps, a sequential cascade only invokes the large model when the small model defers.

In Table 2, we report (i) the reduction in latency from T5 cascades when matching the quality of the large model, and (ii) the best quality that each method can deliver without exceeding the latency of the large model. SpecCascade [Token] often yields the highest speed-up and the best quality metrics, with OPT coming in second. The cascading approaches are often seen to fare poorly on both quality and latency metrics, with the exception of WMT, where SeqCascade yields non-trivial speed-ups. The reason the Token-specific rule fares better than OPT and Diff is because the latter compute their deferral decisions based on which of $q_t(\cdot )$ and $p_t(\cdot )$ is more peaked; this can be a disadvantage when the sampled token is not close to the distribution mode, which is likely to happen with higher temperatures. As shown in § E.3, with lower temperatures, the gap between these rules diminishes.

Few-shot Gemma cascades. To evaluate the Gemma model cascades, we use few-shot prompting with 8 language benchmarks: WMT, CNN/DM, GSM8K, MBPP, SQuAD 2.0, WebQuestions, NaturalQA and TriviaQA; many of these feature in the SpecBench suite ([56]). Figure 3 presents plots of quality vs. rejection rate with a 2B drafter and 27B verifier for $\gamma =1$. For brevity, we only compare the methods that fare the best in the previous experiments. With the exception of TriviaQA, SpecCascade [Token] is able to both match the 27B's quality at a lower rejection rate and yield the best overall quality, often better than 27B. Since all three methods use the exact same implementation for speculative execution, a lower rejection rate directly translates to a lower latency.

Interestingly, OPT is not as effective as with T5. We attribute this to the differences in distributions between the two setups: with T5, the maximum token probability served as a good indicator of token accuracy for both $q$ and $p$; with Gemma, we expect the large model to have a closer alignment with the ground-truth distribution (due to it being several billion parameters apart from the smaller model), and hence using the large model probabilities to measure confidence for both the small and large model (Equation 15) yields better trade-offs than comparing the modes from the two model distributions.

We have proposed new speculative cascading techniques that use a combination of auto-regressive drafting and parallel verification to implement their deferral rule, and shown that they yield better cost-quality trade-offs than standard cascades and speculative decoding. A limitation of our approach is that while it offers a higher throughput, it also incurs a higher total compute cost compared to sequential cascades. In the future, we wish to replace our plug-in estimators with a router model ([25]) trained on ground-truth samples, to improve the local deferral objective at each position $t$ (Equation 8) with a global objective, and to extend our proposal to more than two models.

Proof: Expanding the loss in Equation 3, we have:

This objective is minimized by a deferral rule $r:{\mathcal{V}}^{t-1}\to \{0,1\}$ that minimizes, for each prefix $x_{<t}$, the term within the parenthesis. Therefore the minimizer $r^{\ast }(x_{<t})=1$ whenever the term within the parenthesis is negative:

and $r^{\ast }(x_{<t})=0$ otherwise. Re-arranging the terms completes the proof.

Proof: The proof follows straight-forwardly from the results in ([32]). Recall from § 2 that the lossy speculative decoding procedure of ([32]) accepts a draft token $x$ with probability:

and replaces a rejected draft token with a token sampled from the residual distribution:

for parameters $\alpha \in [0,1)$ and $\beta \ge 1-\alpha$.

We need to show that running Algorithm 4 with the target distribution:

results in the same acceptance probability Equation 16 and residual distribution Equation 17.

The acceptance probability for a draft token $x$ when running Algorithm 4 on $\pi$ is given by:

The corresponding residual distribution is given by:

We consider three possible cases:

Case (i): $q(x)>\frac{1}{1-\alpha }\cdot p(x)\ge \frac{1}{\beta }\cdot p(x)$. In this case, $\pi (x)=\frac{1}{1-\alpha }\cdot p(x)$. As a result:

Case (ii): $\frac{1}{1-\alpha }\cdot p(x)\ge \frac{1}{\beta }\cdot p(x)>q(x)$. In this case, $\pi (x)=\frac{1}{\beta }\cdot p(x)$. As a result:

Case (iii): $\frac{1}{1-\alpha }\cdot p(x)\ge q(x)\ge \frac{1}{\beta }\cdot p(x)$. In this case, $\pi (x)=q(x)$. As a result:

In all three cases, the acceptance probabilities and residual distributions are identical.

Proof: Under a target distribution ${\pi }_t$, the probability of a draft token drawn from $q_t$ being is rejected is given by ([15]):

Expanding $\pi$, the rejection probability becomes:

When $r(x_{<t})=1$, we have:

When $r(x_{<t})=0$, we have:

as desired.

Proof: Expanding the deferral risk in Equation 8, we have:

This objective is minimized by a deferral rule $r:{\mathcal{V}}^{t-1}\to \{0,1\}$ that minimizes, for each prefix $x_{<t}$, the term within the parenthesis. Therefore the minimizer $r^{\ast }(x_{<t})=1$ whenever the term within the parenthesis is negative:

and $r^{\ast }(x_{<t})=0$ otherwise. Re-arranging the terms completes the proof.

For a fixed prefix $x_{<t}$, we can write the deferral risk in Equation 8 as:

where $C$ is a term independent of the deferral rule $r$. Let $r^{\ast }:{\mathcal{V}}^{t-1}\to \{0,1\}$ denote the optimal deferral rule that minimizes $L_{\mathrm{spec}}$ for any prefix $x_{<t}$. We then have: