# Recursive Speculative Decoding:
Accelerating LLM Inference via Sampling Without Replacement

Wonseok Jeon  Mukul Gagrani  Raghavv Goel  Junyoung Park  Mingu Lee
Christopher Lott

###### Abstract

Speculative decoding is an inference-acceleration method for large language
models (LLMs) where a small language model generates a draft-token sequence
which is further verified by the target LLM in parallel. Recent works have
advanced this method by establishing a draft-token tree, achieving superior
performance over a single-sequence speculative decoding. However, those works
independently generate tokens at each level of the tree, not leveraging the
tree’s entire diversifiability. Besides, their empirical superiority has been
shown for fixed length of sequences, implicitly granting more computational
resource to LLM for the tree-based methods. None of the existing works has
conducted empirical studies with fixed target computational budgets despite
its importance to resource-bounded devices. We present Recursive Speculative
Decoding (RSD), a novel tree-based method that samples draft tokens _without_
replacement and maximizes the diversity of the tree. During RSD’s drafting,
the tree is built by either _Gumbel-Top- $k$ trick_ that draws tokens without
replacement in parallel or _Stochastic Beam Search_ that samples _sequences_
without replacement while early-truncating unlikely draft sequences and
reducing the computational cost of LLM. We empirically evaluate RSD with Llama
2 and OPT models, showing that RSD outperforms the baseline methods,
consistently for fixed draft sequence length and in most cases for fixed
computational budgets at LLM.

Machine Learning, ICML

##  1 Introduction

Large language models (LLMs) (Touvron et al., 2023; Zhang et al., 2022; Brown
et al., 2020; Achiam et al., 2023; Jiang et al., 2023) have gained popularity
due to their outstanding achievements with high-quality text generation, which
has drastically increased demands for faster text generation. However, auto-
regressive nature of LLMs limits text generation to produce a single token at
a time and often suffers from memory-bandwidth bottleneck, which leads to
slower inference (Shazeer, 2019). translation (Xiao et al., 2023).

Speculative decoding (Chen et al., 2023; Leviathan et al., 2023) has emerged
as a solution for LLM inference acceleration by leveraging the innate
parallelizability of the transformer network (Vaswani et al., 2017). This
decoding method utilizes a draft model, i.e., a smaller language model, to
auto-regressively generate a sequence of draft tokens with a significantly
lower cost and latency, followed by the target LLM producing the token-wise
probability distributions in parallel. Rejection sampling then verifies those
draft tokens, recovering the sequence distribution by auto-regressive decoding
with the target model. As speculative decoding uses a single sequence of draft
tokens, one needs to increase the draft-sequence length to better exploit
LLM’s parallelizability. However, the longer draft sequence may slow down the
overall inference in practice due to the computational overhead caused by
additional auto-regressive decoding steps from the draft model, possibly
decelerating the target model process due to the increased number of draft
tokens.

Recent works on tree-based speculative decoding (Sun et al., 2023; Miao et
al., 2023) have achieved better diversity and higher acceptance rate via
multiple draft-token sequences. Despite promising results, their decoding
methods independently sample the draft tokens, often harming the diversity of
the tree when samples overlap. Also, their experiments have been conducted for
the fixed length of draft-token sequences across decoding methods, implicitly
requiring more computational resource to the target model when using tree-
based methods. To the best of our knowledge, no prior work has thoroughly
investigated the performance of single-sequence and tree-based speculative
decoding methods with fixed target computational budget, which has practical
importance for resource-bounded devices.

We propose Recursive Speculative Decoding (RSD), a novel tree-based
speculative decoding algorithm that fully exploits the diversity of the draft-
token tree by using sampling without replacement. We summarize our
contributions as below:

Theoretical contribution. We propose _recursive rejection sampling_ capable of
recovering the target model’s distribution with the sampling-without-
replacement distribution defined by the draft model.
Algorithmic contribution. We present RSD which builds draft-token tree
composed of the tokens _sampled without replacement_. Two tree construction
methods, RSD with Constant branching factors (RSD-C) and RSD with Stochastic
Beam Search (RSD-S) (Kool et al., 2019), are proposed.
Empirical contribution. Two perspectives are considered in our experiments:
(Exp1) _performance for fixed length of draft sequence_ , which is also widely
considered in previous works (Sun et al., 2023; Miao et al., 2023), and (Exp2)
_performance for fixed target computational budget_ , where we compared
methods with given size of the draft-token tree. RSD is shown to outperform
the baselines consistently in (Exp1) and for the majority of experiments in
(Exp2).

##  2 Background

Let us consider a sequence generation problem with a set $\mathcal{X}$ of
tokens. We also assume that there is a target model characterized by its
conditional probability
$q(x_{i+1}|x_{1:i}):=\mathrm{Pr}\\{{X_{i+1}}=x_{i+1}|X_{1:i}=x_{1:i}\\}$,
$i\in\mathbb{N}$ for $x_{1:i}:=(x_{1},...,x_{i})$, where
$X_{1},...,X_{i+1}\in\mathcal{X}$ and $x_{1},...,x_{i+1}\in\mathcal{X}$ are
random tokens and their realizations, respectively. Given an input sequence
$X_{1:t}=x_{1:t}$, we can auto-regressively and randomly sample an output
sequence $X_{t+1:t+i}$ for $i\in\mathbb{N}$, i.e., $X_{t+i+1}\sim
q(\cdot|X_{1:t+i}).$

Speculative decoding. Auto-regressive sampling with modern neural network
accelerators (e.g., GPU/TPU) is known to suffer from the memory-bandwidth
bottleneck (Shazeer, 2019), which prevents us from utilizing the entire
computing power of those accelerators. Speculative decoding (Leviathan et al.,
2023; Chen et al., 2023) addresses such issue by using the target model’s
parallelizability. It introduces a (small) draft model which outputs
$p(\hat{X}_{i+1}|\hat{X}_{1:i}):=\mathrm{Pr}\\{\hat{X}_{i+1}=\hat{x}_{i+1}|\hat{X}_{1:i}=\hat{x}_{1:i}\\},i\in\mathbb{N}$.
Speculative decoding accelerates the inference speed by iteratively conducting
the following steps:
1) Draft token generation: For an input sequence $X_{1:m}=x_{1:m}$ and the
draft sequence length $L$, sample draft tokens $\hat{X}_{n+1}\sim
p(\cdot|X_{1:m},\hat{X}_{m+1:n})$ auto-regressively for $n=m,...,m+L-1$ (where
$\hat{X}_{m+1:m}=\emptyset$).
2) Evaluation with target model: Use the target model to compute
$q(\cdot|X_{1:m},\hat{X}_{m+1:n}),n=m,...,m+L$ in parallel.
3) Verification via rejection sampling: Starting from $n=m+1$ to $m+L$,
sequentially accept the draft token $\hat{X}_{n}$ (i.e., $X_{n}=\hat{X}_{n}$)
with the probability
$\min\\{1,\frac{q(\hat{X}_{n}|X_{1:n-1})}{p(\hat{X}_{n}|X_{1:n-1})}\\}$. If
one of the draft tokens $\hat{X}_{n}$ is rejected, we sample $X_{n}\sim
q_{\mathrm{res}}(\cdot|X_{1:n-1})$, where the residual distribution is defined
by

$\displaystyle q_{\mathrm{res}}(\cdot|\tau)$
$\displaystyle:=\mathrm{Norm}[[q(\cdot|\tau)-p(\cdot|\tau)]^{+}],$

for $[f]^{+}:=\max\\{0,f(\cdot)\\}$ and
$\mathrm{Norm}[f]:=\frac{f}{\sum_{x^{\prime}\in\mathcal{X}}f(x^{\prime})}$. If
all draft tokens are accepted ($X_{n}=\hat{X}_{n}$ for $n=m+1,...,m+L$),
sample an extra token $X_{m+L+1}\sim q(\cdot|X_{1:m+L})$.

Chen et al. (2023) and Leviathan et al. (2023) have shown that the target
distribution can be recovered when rejection sampling is applied.

Tree-based speculative decoding. One can further improve the sequence
generation speed by using multiple draft-token sequences, or equivalently, a
tree of draft tokens.

_SpecTr_ (Sun et al., 2023) is a tree-based speculative decoding algorithm
motivated by the Optimal Transport (OT) (Villani et al., 2009). It generalizes
speculative decoding with $K$ i.i.d. draft tokens $\hat{X}^{(k)}\sim
p,k=1,...,K,$ while recovering the target distribution $q$. To this end, a
$K$-sequential draft selection algorithm ($K$-SEQ) was proposed, where the
algorithm decides whether to accept $K$ draft tokens
$\hat{X}^{(k)},k=1,...,K,$ or not with the probability
$\min\\{1,\frac{q(\hat{X}^{(k)})}{\gamma\cdot
p(\hat{X}^{(k)})}\\},\gamma\in[1,K]$. If all draft tokens are rejected, we use
a token drawn from the residual distribution

$\displaystyle\mathrm{Norm}\left[q-\min\left\\{p,\frac{q}{\gamma}\right\\}\frac{1-(1-\beta_{p,q}(\gamma))^{K}}{\beta_{p,q}(\gamma)}\right]$

for $\beta_{p,q}(\gamma):=\sum_{x\in\mathcal{X}}\min\\{p(x),q(x)/\gamma\\}$.

_SpecInfer_ also used the draft-token tree to speed up the inference _with
multiple draft models_ $p^{(k)},k=1,...,K$ (Miao et al., 2023). During the
inference of SpecInfer, all draft models generate their own draft tokwns
independently and create a draft-token tree collectively through repetetion.
For draft verification, _multi-round rejection sampling_ is used to recover
the target distribution, where we determine whether to accept one of the draft
tokens or not with probability
$\min\\{1,\frac{q^{(k)}(\hat{X}^{(k)})}{p^{(k)}(\hat{X}^{(k)})}\\}$ with
distributions $q^{(1)}:=q$ and
$q^{(k)}:=\mathrm{Norm}\left[[q^{(k-1)}-p^{(k-1)}]^{+}\right],k=2,...,K+1.$ If
all draft tokens are rejected, we sample a token $Y\sim q^{(K+1)}$ from the
last residual distribution.

##  3 Recursive Speculative Decoding

In this section, we present Recursive Speculative Decoding (RSD), a tree-based
speculative decoding method that constructs draft-token trees via sampling
without replacement. We first propose recursive rejection sampling, a
generalization of multi-round speculative decoding (Miao et al., 2023) that is
applicable to draft distributions with dependencies, where sampling-without-
replacement distribution is one instance of such distributions. Then, we use
recursive rejection sampling to validate each level of the draft-token tree
which can be efficiently constructed via either Gumbel-Top-$k$ trick (Vieira,
2014) and Stochastic Beam Search (Kool et al., 2019),

Algorithm 1 Recursive Rejection Sampling

1: Input: Draft dist. $p^{(k)},k=1,...,K,$ target dist. $q$.

2: Sample $\hat{X}^{(k)}$ by (1).

3: Compute $q^{(k)}(\cdot|\hat{X}^{(1:k-2)})$ and $\Theta^{(k)}$ by (2) and
(3).

4: for $k$ in $\\{1,...,K\\}$ do

5: Sample $A^{(k)}\in\\{\mathtt{acc},\mathtt{rej}\\}$ with probability
$\Theta^{(k)}$.

6: if $A^{(k)}=\mathtt{acc}$ then return $Z\leftarrow\hat{X}^{(k)}$; end if

7: end for

8: return $Z\sim q^{(K+1)}(\cdot|\hat{X}^{(1:K-1)})$

###  3.1 Recursive Rejection Sampling: Generalized Multi-Round Rejection
Sampling

Suppose we have target distribution $q(x),x\in\mathcal{X}$. In recursive
rejection sampling, we introduce random variables
$\hat{X}^{(1)},...,\hat{X}^{(K)}\in\mathcal{X}$ that represent $K$ draft
tokens; these tokens will locate at the same level of the draft-token tree in
Section 3.2. We aim to recover target distribution $q$, where

$\displaystyle\hat{X}^{(1)}\sim p^{(1)},\hat{X}^{(k)}\sim
p^{(k)}(\cdot|\hat{X}^{(1:k-1)}),k=2,...,K,$ (1)

for some distributions $p^{(k)},k=1,...,K$ and a sequence
$\hat{X}^{(1:k-1)}:=(\hat{X}^{(1)},...,\hat{X}^{(k-1)})$. Note that we assume
distributions with dependencies unlike prior works such as SpecTr (Sun et al.,
2023) consider independent distributions. By using $p^{(1)},...,p^{(K)}$ and
$q$, we define $q^{(1)}:=q$ and residual distributions

$\displaystyle q^{(k+1)}(\cdot|x^{(1:k-1)})$
$\displaystyle:=\mathrm{Norm}\left[[q^{(k)}(\cdot|x^{(1:k-2)})-p^{(k)}(\cdot|x^{(1:k-1)})]^{+}\right]$
(2)

for $k=1,...,K$ and $x^{(1)},...,x^{(K+1)}\in\mathcal{X}$, where
$x^{(1:k^{\prime})}=\emptyset$ (empty sequence, i.e., no conditioning) if
$k^{\prime}<1$, or $(x^{(1)},...,x^{(k^{\prime})})$, otherwise. Together with
draft, target, and residual distributions, recursive rejection sampling
introduces threshold random variables $\Theta^{(1)},...,\Theta^{(K)}\in[0,1]$
which determines rejection criteria for each draft token
$\hat{X}^{(k)},k=1,...,K$:

$\displaystyle\Theta^{(k)}$
$\displaystyle:=\min\left\\{1,\frac{q^{(k)}(\hat{X}^{(k)}|\hat{X}^{(1:k-2)})}{p^{(k)}(\hat{X}^{(k)}|\hat{X}^{(1:k-1)})}\right\\}.$
(3)

Specifically, each $\Theta^{(k)}$ can be used to define random variables
$A^{(k)}\in\\{\mathtt{acc},\mathtt{rej}\\}$ (where $\mathtt{acc}$ and
$\mathtt{rej}$ indicate acceptance and rejection of draft tokens,
respectively) such that
$\mathrm{Pr}\left\\{A^{(k)}=\mathtt{acc}|\Theta^{(k)}=\theta\right\\}=\theta$
for $\theta\in[0,1]$.

Figure 1:  Acceptance rates for multi-round speculative decoding, K-SEQ, OTM
and recursive rejection sampling are given when $\mathrm{Ber}(p)$ and
$\mathrm{Ber}(q)$ are draft and target distributions, respectively, and two
tokens are proposed by the draft model ($K=2$).

Figure 2: We describe the entire process of RSD with Stochastic Beam Search
(RSD-S); the difference between RSD-S and RSD with Constant branching factors
(RSD-C) lies at the method of constructing the draft-token tree. Draft tokens
the tree are sampled in parallel at each level and auto-regressively across
levels, while Stochastic Beam Search samples sequences without replacement at
each tree level. The established draft-token tree is then processed by the
target model in parallel, which lets us acquire the token-wise target model
probabilities. Finally, recursive rejection sampling (for sampling-without-
replacement distribution) is applied to each level of the tree, recovering the
sequence generation distribution of the target model.

Finally, recursive rejection sampling can be characterized by defining a
random variable $Z\in\mathcal{X}$ such that

$\displaystyle
Z:=\begin{cases}\hat{X}^{(k)},&\text{if~{}}A^{(1:k-1)}=\mathtt{rej}^{k-1},A^{(k)}=\mathtt{acc},\\\
&{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\text{if~{}}}k=1,...,K,\\\
Y,&\text{if~{}}A^{(1:K)}=\mathtt{rej}^{K},\\\
&{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\text{if~{}}}Y\sim
q^{(K+1)}(\cdot|\hat{X}^{(1:K-1)}),\end{cases}$ (4)

where $A^{(1:k-1)}:=(A^{(1)},...,A^{(k-1)})$ and $\mathtt{rej}^{k}$ is a
length-$k$ sequence with all of its elements equal to $\mathtt{rej}$.
Intuitively, we select $\hat{X}^{(1)}$ if it is accepted
($A^{(1)}=\mathtt{acc}$); we select $\hat{X}^{(k)}$ when all previous draft
tokens $\hat{X}^{(1)},...,\hat{X}^{(k-1)}$ are rejected _and_ $\hat{X}^{(k)}$
is accepted ($A^{(1:k-1)}=\mathtt{rej}^{k-1},A^{(k)}=\mathtt{acc}$) for each
$k$; we sample $Y\sim q^{(K+1)}(\cdot|\hat{X}^{(1:K-1)})$ and select $Y$ if
all draft tokens are rejected ($A^{(1:K)}=\mathtt{rej}^{K}$). We summarize the
entire process of recursive rejection sampling in Algorithm 1. Note that the
original rejection sampling (Leviathan et al., 2023; Chen et al., 2023) is a
special case of our recursive rejection sampling with $K=1$. Also, it can be
shown that recursive rejection sampling (4) always recovers the target
distribution $q$:

######  Theorem 3.1 (Recursive rejection sampling recovers target
distribution).

For the random variable $Z\in\mathcal{X}$ in (4),

$\displaystyle\mathrm{Pr}\\{Z=z\\}=q(z),z\in\mathcal{X}.$

###### Proof.

See Appendix A.1. ∎

Although the proposed recursive rejection sampling is applicable to arbitrary
distributions with dependencies following (1), we assume a single draft model
(as in SpecTr (Sun et al., 2023) and focus on the cases where the draft model
samples predictive tokens without replacement, which is an instance of (1).

Toy example. We present a didactic example with Bernoulli distributions (given
by Sun et al. (2023)) to showcase the benefit of recursive rejection sampling.
Suppose that Bernoulli distributions are used for both draft and target models
and only $K=2$ tokens are allowed for draft proposals. The acceptance rates
for different methods are depicted in Figure 1; multi-round speculative
decoding (from SpecInfer (Miao et al., 2023)), K-SEQ and Optimal Transport
with Membership costs (OTM) (Sun et al., 2023), use sampling _with_
replacement, whereas recursive rejection sampling uses sampling _without_
replacement; note that both K-SEQ and OTM were presented in SpecTr paper (Sun
et al., 2023) where OTM shows theoretically optimal acceptance rate. For all
the baselines, acceptance rates decrease as the discrepancy between draft and
target distribution increases, since tokens sampled from draft models become
more unlikely from target models. On the other hand, recursive rejection
sampling achieves 100% acceptance rate even with high draft-target-model
discrepancy; once the first draft token is rejected, the second draft token is
always aligned with the residual distribution. This example shows that draft
distributions with dependencies, e.g., sampling-without-replacement
distribution, leads to higher acceptance rate and becomes crucial, especially
for the cases with higher distributional discrepancy between draft and target.

###  3.2 Tree-Based Speculative Decoding with Recursive Rejection Sampling

Recursive rejection sampling is applicable to tree-based speculative decoding
algorithms if sampling without replacement is used to construct a _draft-token
tree_. Two Recursive Speculative Decoding (RSD) algorithms using recursive
rejection sampling are presented in this section, while they share the same
pipeline for parallel target evaluation and draft tree verification after
building the draft-token tree (See _Figure 2._). We describe details about how
RSD works in the following sections.

####  3.2.1 Draft-token tree generation

We consider two RSD algorithms: RSD with Constant branching factors (RSD-C)
and RSD with Stochastic Beam Search (RSD-S). RSD-C builds the draft-token tree
having constant branching factors, which makes sequences from the tree to have
the same length. RSD-S, on the other hand, builds the tree via Stochastic Beam
Search (Kool et al., 2019) that samples _draft sequences_ without replacement,
while truncating sequences that are unlikely to be generated from the draft
model and efficiently handling the computational cost.

RSD with Constant branching factors (RSD-C). Let $L$ denote the fixed length
for all draft sequences, which is equivalent to the depth of the draft-token
tree, and $\boldsymbol{\tau}_{0}^{(1)}$ denote the input sequence of tokens.
Let us assume that the tree level increases from root ($l=0$) to leaf ($l=L$)
nodes, where each node is characterized by the (partial) sequence. We also
define $\mathbf{b}:=(b_{0},...,b_{L-1})$ where $b_{l}$ is the branching factor
at the level $l$ (See Figure 3(a) for the example with $\mathbf{b}=(3,2,1)$.).

At each level $l\in\\{0,...,L-1\\}$ of the draft tree, we begin with $N_{l}$
sequences $\boldsymbol{\tau}_{l}^{(k)},k=1,...,N_{l}$ generated from the
previous level, where $N_{0}:=1$ and
$N_{l}:=\prod_{l^{\prime}=0}^{l-1}b_{l^{\prime}}$ for $l\geq 1$. Then, we
evaluate log probabilities $\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ and
perturbed log probabilities
$\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ for each $k$, i.e., for
i.i.d. Gumbel samples $G_{l}^{(k)}(x),x\in\mathcal{X},$

$\displaystyle\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$
$\displaystyle:=\log p(\cdot|\boldsymbol{\tau}_{l}^{(k)}),$ (5)
$\displaystyle\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$
$\displaystyle:=\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)+G_{l}^{(k)},$ (6)

where both log probabilities and Gumbel samples can be computed in parallel;
proper positional encodings and attention masking (Cai et al., 2023; Miao et
al., 2023) are required for the parallel log-probability computation when
transformer architecture is used (Vaswani et al., 2017). By using _Gumbel-Top-
$k$ trick_ (Vieira, 2014; Kool et al., 2019) with perturbed log probabilities
(6), one can sample top-$b_{l}$ tokens without replacement for each sequence
$\tau_{l}^{(k)}$:

$\displaystyle\hat{X}_{l+1}^{((k-1)b_{l}+1)},...,\hat{X}_{l+1}^{((k-1)b_{l}+b_{l})}=\underset{x\in\mathcal{X}}{\mathrm{argtop}\text{-}b_{l}}\left(\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},x)\right).$
(7)

Note that the outputs
$\hat{X}_{l+1}^{((k-1)b_{l}+k^{\prime})},k^{\prime}=1,...,b_{l},$ in (7) are
assumed to be in the decreasing order of values
$\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\hat{X}_{l+1}^{((k-1)b_{l}+k^{\prime})})$,
for each $k$. Finally, we define

$\displaystyle O_{l+1}^{((k-1)b_{l}+k^{\prime})}$
$\displaystyle:=(\hat{X}_{l+1}^{((k-1)b_{l}+k^{\prime})},k),$ (8)
$\displaystyle\boldsymbol{\tau}_{l+1}^{((k-1)b_{l}+1)}$
$\displaystyle:=(\boldsymbol{\tau}_{l}^{(k)},\hat{X}_{l+1}^{((k-1)b_{l}+1)})$
(9)

for $k\in{1,...,N_{l}}$ and $k^{\prime}\in\\{1,...,b_{l}\\}$, where
$O_{l+1}^{((k-1)b_{l}+k^{\prime})}$ is a pair of draft token and parent
sequence index. Those pairs in (8) are stored for all levels $l=0,...,L-1$ and
used for draft tree verification, which exploits the fact that _the tokens
$\hat{X}_{l+1}^{((k-1)b_{l}+1)},...,\hat{X}_{l+1}^{((k-1)b_{l}+b_{l})}$ follow
sampling without replacement from $p(\cdot|\boldsymbol{\tau}_{l}^{(k)})$ for
any given parent sequence index $k$_.

(a) RSD-C(b) RSD-S Figure 3: We describe examples of constructing draft-token
trees with the (maximum) draft length equal to 3; (a) The tree constructed by
RSD-C with branching factors $\mathbf{b}=(3,2,1)$ is given; (b) we depict the
tree constructed by RSD-S with beamwidth $W=3$, where edges are determined via
Stochastic Beam Search.

RSD with Stochastic Beam Search (RSD-S). One caveat of RSD-C is that its
constant branching factors $\mathbf{b}$ should be carefully determined to
handle tree complexity, when the computation budget is limited; for example,
if $\mathbf{b}=(n,...,n)$ with its length $L$, the number of nodes in the
draft tree will be $\sum_{l=0}^{L-1}n^{l}=O(n^{L-1})$, which is
computationally prohibitive for large $n$ and $L$. Also, RSD-C constructs
sequences at each level $l$ by using the _myopic_ token-wise log probabilities
$\phi_{l}$ in (6). RSD-S addresses both issues by using Stochastic Beam Search
(Kool et al., 2019) that early-truncates unlikely sequences and utilizes _far-
sighted_ sequence log probabilities.

Let us define the _maximum_ draft sequence length $L$ and the beamwidth $W$.
We also define $\boldsymbol{\tau}_{0}^{(1)}$ as the input sequence similar to
RSD-C. At each level $l\in\\{0,...,L-1\\}$, SBS uses _beam_

$\displaystyle\mathcal{B}_{l}$ $\displaystyle:=(t_{l}^{(1)},...,t_{l}^{(W)}),$
$\displaystyle t_{l}^{(k)}$
$\displaystyle:=(\boldsymbol{\tau}_{l}^{(k)},\phi_{l-1}(\boldsymbol{\tau}_{l}^{(k)}),\psi_{l-1}(\boldsymbol{\tau}_{l}^{(k)}))$

generated from the previous level $l-1$111For $l=0$,
$\phi_{-1}(\boldsymbol{\tau}_{0}^{(1)})=\phi_{-1}(\boldsymbol{\tau}_{0}^{(1)})=0$
is used with $\mathcal{B}_{0}:=(t_{0}^{(1)})$ (Kool et al., 2019).. Here, each
tuple $t_{l}^{(k)}$ for $k\in\\{1,...,W\\}$ consists of _(a)_ a sequence
$\boldsymbol{\tau}_{l}^{(k)}$, _(b)_ its _sequence_ log probability
$\phi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})$ of $\boldsymbol{\tau}_{l}^{(k)}$,
and _(c)_ the _transformed (perturbed and truncated)_ sequence log probability
$\psi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})$, respectively.

For each tuple $t_{l}^{(k)}$ in the beam $\mathcal{B}_{l}$, we evaluate the
(next-level) sequence log probabilities
$\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ and the perturbed sequence log
probabilities $\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$.
Specifically for i.i.d. Gumbel samples $G_{l}^{(k)}(x),x\in\mathcal{X}$, we
compute

$\displaystyle\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$
$\displaystyle:={\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\phi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})}+\log
p(\cdot|{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\boldsymbol{\tau}_{l}^{(k)}}),$
$\displaystyle\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$
$\displaystyle:=\phi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)+G_{l}^{(k)},$

where the terms $\boldsymbol{\tau}_{l}^{(k)}$ and
$\phi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})$ within the tuple $t_{l}^{(k)}$ of
within the beam $\mathcal{B}_{l}$ are reused. Similar to RSD-C, both log
probabilities and Gumbel samples can be parallelly computed with positional
encodings and attention masking (Cai et al., 2023; Miao et al., 2023). In
addition to the perturbed log probabilities, SBS in RSD-S transforms
$\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ into the truncated
function

$\displaystyle\psi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$
$\displaystyle:=T({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\psi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})},\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)),$
(10) $\displaystyle T(u,\phi)$
$\displaystyle:=-\log\left(e^{-u}-e^{-\mathrm{max}\phi}+e^{-\phi(\cdot)}\right)$
(11)

for $\max\phi:=\max_{x\in\mathcal{X}}\phi(x)$ by reusing
$\psi_{l-1}(\boldsymbol{\tau}_{l}^{(k)})$ in $t_{l}^{(k)}$ . Note that
$T(u,\phi)$ in (11) is _monotonically increasing_ w.r.t. $\phi$ and transforms
$\phi$ to the function with the upper bound $u$ (Kool et al., 2019)222In
Appendix B.3 of Kool et al. (2019), a numerical stable way of evaluating the
function $T$ in (11) is provided.

After evaluating $\psi_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ for all parent
sequences $\boldsymbol{\tau}_{l}^{(k)}$s, SBS selects top-$W$ pairs
$(\hat{X}_{l+1},p_{l+1})$ of draft token and parent sequence index _across the
beam $\mathcal{B}_{l}$_, i.e.,

$\displaystyle
O_{l+1}^{(1)},...,O_{l+1}^{(W)}:=\underset{(x,k)\in\mathcal{X}\times\mathcal{K}}{\mathrm{argtop}\text{-}W}\left(\psi_{l}(\boldsymbol{\tau}_{l}^{(k)},x)\right)$
(12)

for $O_{l+1}^{(k)}:=(\hat{X}_{l+1}^{(k)},p_{l+1}^{(k)})$ and
$\mathcal{K}:=\\{1,...,W\\}$. The output pairs
$O_{l+1}^{(1)},...,O_{l+1}^{(W)}$ are given by _corresponding values
$\psi_{l}(\boldsymbol{\tau}_{l}^{(k)},\hat{X}_{l+1}^{(k)})$ in the decreasing
order_. Finally, we construct the next beam

$\displaystyle\mathcal{B}_{l+1}$
$\displaystyle:=(t_{l+1}^{(1)},...,t_{l+1}^{(W)}),$ $\displaystyle
t_{l+1}^{(k)}$
$\displaystyle:=((\hat{\boldsymbol{\tau}}_{l+1}^{(k)},\hat{X}_{l+1}^{(k)}),\phi_{l}(\hat{\boldsymbol{\tau}}_{l+1}^{(k)},\hat{X}_{l+1}^{(k)}),\psi_{l}(\hat{\boldsymbol{\tau}}_{l+1}^{(k)},\hat{X}_{l+1}^{(k)}))$

for $k=1,...,W$, where
$\hat{\boldsymbol{\tau}}_{l+1}^{(k)}:=\boldsymbol{\tau}_{l}^{(p_{l+1}^{(k)})}$
is the selected parent sequence. Intuitively, SBS at the level $l$ evaluates
scores
$\psi_{l}^{(k)}(\boldsymbol{\tau}_{l}^{(k)},x),x\in\mathcal{X},k\in\mathcal{K},$
by considering all child nodes from the beam $\mathcal{B}_{l}$. SBS selects
$W$ nodes among all child nodes having top-$W$ scores. Note that the above
process is theoretically equivalent to sample top-$W$ length-$(l+1)$ sequences
_without replacement_ (Kool et al., 2019) and efficiently truncates sequences
that are unlikely to be generated. (See _Figure 3(b)_.)

We store the _ordered_ sequence of pairs $O_{l+1}^{(1)},...,O_{l+1}^{(W)}$ for
all levels $l=0,...,L-1$, which is used for draft-tree verification. As in
RSD-C, we show the following property:

######  Theorem 3.2 (Tokens from the same sequence follow sampling without
replacement in RSD-S).

In RSD-S, any non-empty subsequence of the sequence
$\hat{X}_{l+1}^{(1)},...,\hat{X}_{l+1}^{(W)}$ of draft tokens (from
$O_{l+1}^{(1)},...,O_{l+1}^{(W)}$ in (12)) such that each element of the
subsequence has the same parent $\boldsymbol{\tau}_{l}^{(k)}$ follows sampling
without replacement from $p(\cdot|\boldsymbol{\tau}_{l}^{(k)})$333We define a
subsequence of a sequence as any sequence acquired by removing its elements
_while maintaining the order in the original sequence_..

###### Proof.

See Appendix A.2. ∎

####  3.2.2 Draft-tree evaluation and verification

Tree evaluation with target model. After the draft-tree construction, we have
sequences of pairs

$\displaystyle(O_{l}^{(1)},...,O_{l}^{(N_{l})}),l=1,...,L,$

where $N_{l}=\prod_{l^{\prime}=0}^{l}b_{l^{\prime}}$ for RSD-C and $N_{l}=W$
for RSD-S, respectively ($N_{0}:=1$ for both). Those pairs include the node-
connection information of the draft tree and can be used to _parallelly_
evaluate the draft tree via the target model by utilizing appropriate
attention masking and positional encodings. From the evaluation process, we
acquire the target log probabilities for all sequences
$\boldsymbol{\tau}_{l}^{(k_{l})}$ in the draft tree, i.e.,

$\displaystyle
q(\cdot|\boldsymbol{\tau}_{l}^{(k_{l})}),l=0,...,L,k_{l}=1,...,N_{l}.$

Verification via recursive rejection sampling. Earlier, we show that tokens in
the tree having the same parent sequence $\boldsymbol{\tau}_{l}^{(k_{l})}$
follows the sampling-without-replacement distribution from
$p(\cdot|\boldsymbol{\tau}_{l}^{(k_{l})})$ for both RSD-C and RSD-S. Thus, one
can apply recursive rejection sampling iteratively at each tree level.

Specifically, at the level $l\in\\{0,1,...,L\\}$, we begin with a sequence
$\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})}$ where $k_{l}^{\prime}$ is the index
of the parent sequence accepted in the previous level ($k_{0}^{\prime}=1$ at
the level $l=0$). Within the _ordered_ sequence
$(O_{l+1}^{(1)},...,O_{l+1}^{(N_{l+1})})$ of pairs, we find the subsequence
$\mathbf{o}_{l+1}^{(k_{l}^{\prime})}$ having
$\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})}$ as parent, which can be validated
by checking the second element of each pair $O_{l+1}^{(k)}$, and the token
sequence $\mathbf{x}_{l+1}^{(k_{l}^{\prime})}$ in
$\mathbf{o}_{l+1}^{(k_{l}^{\prime})}$. Earlier, we show that tokens
$\mathbf{x}_{l+1}^{(k_{l}^{\prime})}$ follows sampling-without-replacement
distribution in its order, so we can apply recursive rejection sampling to
those tokens with draft and target distributions,
$p(\cdot|\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})})$ and
$q(\cdot|\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})})$, respectively. If any
token $x$ in $\mathbf{x}_{l+1}^{(k_{l}^{\prime})}$ is accepted, we set
$k_{l+1}^{\prime}$ that corresponds to
$\boldsymbol{\tau}_{l}^{(k_{l+1}^{\prime})}:=(\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})},x)$,
and we continue to the next-level verification if child nodes exist. If all
tokens are rejected, we sample from the last residual distribution of
recursiver rejection sampling. If there is no child node, we sample from the
target $q(\cdot|\boldsymbol{\tau}_{l}^{(k_{l}^{\prime})})$ similar to the
single-sequence speculative decoding (Chen et al., 2023; Leviathan et al.,
2023). We provide detailed descriptions for RSD-C (Algorithm 2) and for RSD-S
(Algorithm 7) in Appendix B.

Llama 2-70BLlama 2-Chat-70B Figure 4: Block efficiency, MBSU, token rate and
accuracy for various lengths ($2,3,4,5$) of draft sequences are given. We
consider two target models, Llama 2-70B and Llama 2-Chat-70B, each of which
has a corresponding smaller draft model for speculative decoding. All results
are normalized by the corresponding numbers from auto-regressive decoding.
RSD-S always outperforms SD, SpecTr and RSD-C. All methods including auto-
regressive decoding show similar accuracy for WMT and XSum.  Llama 2-70BLlama
2-Chat-70B Figure 5: Block efficiency, MBSU, token rate and accuracy for
various target computational budgets (the numbers $6,10,14,21,30$ of draft
tokens processed at the target model) are given. We consider two target
models, Llama 2-70B and Llama 2-Chat-70B, each of which has a corresponding
smaller draft model for speculative decoding. All results are normalized by
the corresponding numbers from auto-regressive decoding. RSD-S outperforms SD,
SpecTr and RSD-C in the majority of cases. All methods including auto-
regressive decoding show similar accuracy for both WMT and XSum.

##  4 Related Works

Many recent works have aimed to address the inference bottleneck of LLMs
caused by auto-regressive decoding. Speculative decoding methods (Leviathan et
al., 2023; Chen et al., 2023; Sun et al., 2023; Miao et al., 2023) use the
target model (LLM) with a draft model (a small language model), while
recovering target distribution via rejection sampling. See the recent survey
on speculative decoding (Xia et al., 2024) for more comprehensive
understanding.

Other than speculative decoding methods, BiLD (Kim et al., 2023) is another
method to accelerate inference, where it uses a fallback policy which
determines when to invoke the target model and a rollback policy to review and
correct draft tokens. Medusa (Cai et al., 2024) uses multiple decoding heads
to predict future tokens in parallel, constructs the draft-token tree and uses
a typical acceptance criteria. Lookahead decoding (Fu et al., 2023) caches the
historical $n$-grams generated on-the-fly instead of having a draft model and
performs parallel decoding using Jacobi iteration and verifies $n$-grams from
the cache. While showing promising results with greedy sampling, these works
do not guarantee target distribution recovery in contrast to speculative
decoding methods.

##  5 Experiments

We evaluate RSD-C and RSD-S together with our baselines including speculative
decoding (SD) (Chen et al., 2023; Leviathan et al., 2023) and SpecTr (Sun et
al., 2023), where a single draft model is assumed444We exclude SpecInfer (Miao
et al., 2023) from our baselines since it uses multiple draft models.. We
consider the following perspectives during our experiments: (Exp1) How will
the performance be affected by _the length of draft sequences_? (Exp2) How
will the performance be affected by _the target computational budget_ , i.e.,
the number of tokens processed at the target model? While (Exp1) has been
frequently investigated by existing tree-based speculative decoding methods
(Sun et al., 2023; Miao et al., 2023), (Exp2) has not been considered in prior
works but has practical importance when running the target model on resource-
bounded devices.
Models. We consider the following target models; Llama 2 and Llama 2-Chat
(Touvron et al., 2023) with 7B, 13B and 70B parameters; OPT (Zhang et al.,
2022) with 13B, 30B and 66B parameters. Each class of target models adopts
corresponding draft model; see Appendix C.1. In this section, we only present
Llama 2-70B and Llama 2-Chat-70B results, and other results (Llama 2 with
other sizes and OPT) can be found in Appendix C.4.
Tasks. Our methods and baselines are evaluated for WMT18-DeEn (Bojar et al.,
2018, translation) and XSum (Narayan et al., 2018, summarization) for each
target model, while we report accuracy scores (BLEU for WMT and ROUGE-2 for
XSum) to confirm if the target model’s distribution is recovered; Databricks-
Dolly-15k (Conover et al., 2023, question and answering) is used only for
Llama 2-Chat without accuracy evaluation. We use temperature 0.3 for both XSum
and WMT and 1.0 for Dolly, where we further apply nucleus (top-$p$) sampling
(Holtzman et al., 2019) with $p=0.95$ for Dolly.
Performance metrics. We evaluate block efficiency (Leviathan et al., 2023),
Memory-Bound Speed-Up (MBSU) (Zhou et al., 2023) and token rate (tokens/sec)
on A100 GPUs; see Appendix C.2 for details.

###  5.1 (Exp 1) Fixed draft sequence length

We fix (maximum) draft sequence length as the value in $\\{2,3,4,5\\}$ and
evaluate our methods and baselines, which is summarized in _Figure 4_.
Regarding the tree structures of each decoding methods, we let both SpecTr and
RSD-S always use draft-token trees, the size of which is smaller than or equal
to that of RSD-C’s tree; see Appendix C.3.1 for details. Our results show that
tree-based methods (SpecTr, RSD-C and RSD-S) always outperform SD in terms of
block efficiency and MBSU, whereas token rates for SpecTr and RSD-C can be
lower than that for SD; this is since block efficiencies for both SpecTr and
RSD-C are relatively low and there is additional computational overhead to
process the tree. On the other hand, _RSD-S strictly outperforms both SD and
SpecTr for all performance metrics_ , showing the superiority of RSD-S over
our baselines and the importance of early-truncating unlikely draft sequences.
We also observe that there is no strong correlation between MBSU and token
rate; this is since A100 GPUs used to measure token rates are _not_ memory-
bound. Furthermore, token rates in many cases are shown to decrease as the
length of draft-token sequence becomes higher, which is due to the increased
computation overhead to execute draft models with the longer draft sequence;
however, one needs to be cautious since this result may not generally hold
since token rate is hugely affected by the efficiency of software
implementation and the devices which we execute the methods on. Finally, in
WMT and XSum, BLEU and ROUGE-2 scores are similar across different methods,
respectively, which implies that all methods recover the distributions of
target LLMs.

###  5.2 (Exp2) Fixed target computational budget

We select target computational budget, i.e., the number of draft tokens
processed at the target model in parallel for each speculative decoding
iteration, among values in $\\{6,10,14,21,30\\}$ and evaluate our proposed
methods and baselines; we summarize the results in _Figure 5_ and describe
tree structures in Appendix C.3.2. While RSD-S achieves higher block
efficiency and MBSU than SD and SpecTr in most cases, SD beats RSD-C in the
relatively low budget regime $\\{6,10,14\\}$ in XSum experiments, when using
Llama 2-70B as a target model. We believe that our draft model is well-aligned
with Llama 2-70B in XSum task (from the observation that block efficiency
between 2.4 and 3.0, which is significantly higher than the numbers in other
cases, is achieved), and increasing the depth rather than the width of the
tree could quickly increase the acceptance rate in this case. In the high
budget regime ($\\{21,30\\}$), on the other hand, RSD-S beats SD for both
block efficiency and MBSU. In terms of token rate, RSD-S strictly outperforms
our baselines, whereas SD’s token rate severely decreases for higher target
computation budgets due to the computational overhead caused by the draft
model’s auto-regressive decoding with the longer draft sequence.

##  6 Conclusion

We present RSD algorithms, a novel tree-based speculative decoding method
leveraging the full diversifiability of the draft-token tree; RSD-C
efficiently samples draft tokens without replacement via Gumbel-Top-$k$ trick,
while RSD-S uses Stochastic Beam Search and samples draft-token sequences
without replacement. We also propose recursive rejection sampling that can
verify the tree built by the sampling-without-replacement process and recovers
the exact target model distribution. We show that RSD outperforms the
baselines in most cases, supporting the importance of diverse drafting when
accelerating LLM inference.

## References

  * Achiam et al. (2023) Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al.  GPT-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Bojar et al. (2018) Bojar, O. r., Federmann, C., Fishel, M., Graham, Y., Haddow, B., Huck, M., Koehn, P., and Monz, C.  Findings of the 2018 conference on machine translation (wmt18).  In _Proceedings of the Third Conference on Machine Translation, Volume 2: Shared Task Papers_ , pp. 272–307, Belgium, Brussels, October 2018. Association for Computational Linguistics.  URL http://www.aclweb.org/anthology/W18-6401.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 33:1877–1901, 2020.
  * Cai et al. (2023) Cai, T., Li, Y., Geng, Z., Peng, H., and Dao, T.  Medusa: Simple framework for accelerating llm generation with multiple decoding heads.  https://github.com/FasterDecoding/Medusa, 2023.
  * Cai et al. (2024) Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T.  Medusa: Simple llm inference acceleration framework with multiple decoding heads.  _arXiv preprint arXiv:2401.10774_ , 2024.
  * Chen et al. (2023) Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J.  Accelerating large language model decoding with speculative sampling.  _arXiv preprint arXiv:2302.01318_ , 2023.
  * Conover et al. (2023) Conover, M., Hayes, M., Mathur, A., Xie, J., Wan, J., Shah, S., Ghodsi, A., Wendell, P., Zaharia, M., and Xin, R.  Free dolly: Introducing the world’s first truly open instruction-tuned llm, 2023.  URL https://www.databricks.com/blog/2023/04/12/dolly-first-open-commercially-viable-instruction-tuned-llm.
  * Fu et al. (2023) Fu, Y., Bailis, P., Stoica, I., and Zhang, H.  Breaking the sequential dependency of LLM inference using lookahead decoding, November 2023.  URL https://lmsys.org/blog/2023-11-21-lookahead-decoding/.
  * Holtzman et al. (2019) Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y.  The curious case of neural text degeneration.  _arXiv preprint arXiv:1904.09751_ , 2019.
  * Jiang et al. (2023) Jiang, A. Q., Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. l., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al.  Mistral 7B.  _arXiv preprint arXiv:2310.06825_ , 2023.
  * Kim et al. (2023) Kim, S., Mangalam, K., Malik, J., Mahoney, M. W., Gholami, A., and Keutzer, K.  Big little transformer decoder.  _arXiv preprint arXiv:2302.07863_ , 2023.
  * Kool et al. (2019) Kool, W., Van Hoof, H., and Welling, M.  Stochastic beams and where to find them: The Gumbel-Top-$k$ trick for sampling sequences without replacement.  In _Proceedings of the 36th International Conference on Machine Learning (ICML)_ , pp. 3499–3508. PMLR, 2019.
  * Leviathan et al. (2023) Leviathan, Y., Kalman, M., and Matias, Y.  Fast inference from transformers via speculative decoding.  In _Proceedings of the 40th International Conference on Machine Learning (ICML)_ , 2023.
  * Miao et al. (2023) Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z.  SpecInfer: Accelerating generative LLM serving with speculative inference and token tree verification.  _arXiv preprint arXiv:2305.09781_ , 2023.
  * Narayan et al. (2018) Narayan, S., Cohen, S. B., and Lapata, M.  Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization.  _arXiv preprint arXiv:1808.08745_ , 2018.
  * Shazeer (2019) Shazeer, N.  Fast transformer decoding: One write-head is all you need.  _arXiv preprint arXiv:1911.02150_ , 2019.
  * Sun et al. (2023) Sun, Z., Suresh, A. T., Ro, J. H., Beirami, A., Jain, H., and Yu, F.  SpecTr: Fast speculative decoding via optimal transport.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2023.
  * Touvron et al. (2023) Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_ , 2023.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 30, 2017.
  * Vieira (2014) Vieira, T.  Gumbel-max trick and weighted reservoir sampling.  2014\.  URL https://timvieira.github.io/blog/post/2014/08/01/gumbel-max-trick-andweighted-reservoir-sampling/.
  * Villani et al. (2009) Villani, C. et al.  _Optimal transport: old and new_ , volume 338.  Springer, 2009.
  * Xia et al. (2024) Xia, H., Yang, Z., Dong, Q., Wang, P., Li, Y., Ge, T., Liu, T., Li, W., and Sui, Z.  Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding.  _arXiv preprint arXiv:2401.07851_ , 2024.
  * Xiao et al. (2023) Xiao, Y., Wu, L., Guo, J., Li, J., Zhang, M., Qin, T., and Liu, T.-y.  A survey on non-autoregressive generation for neural machine translation and beyond.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 2023.
  * Zhang et al. (2022) Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al.  OPT: Open pre-trained transformer language models.  _arXiv preprint arXiv:2205.01068_ , 2022.
  * Zhou et al. (2023) Zhou, Y., Lyu, K., Rawat, A. S., Menon, A. K., Rostamizadeh, A., Kumar, S., Kagy, J.-F., and Agarwal, R.  Distillspec: Improving speculative decoding via knowledge distillation.  _arXiv preprint arXiv:2310.08461_ , 2023.

##  Appendix A Theorems and proofs

###  A.1 Proof of Theorem 3.1

######  Theorem 3.1 (Recursive rejection sampling recovers target
distribution).

The random variable $Z\in\mathcal{X}$ defining recursive rejection sampling
rule (4) follows the target distribution $q$, i.e.,

$\displaystyle\mathrm{Pr}\left\\{Z=z\right\\}=q(z),z\in\mathcal{X}.$

###### Proof.

We remain a sketch of the proof here and the formal proof is given in the next
paragraph. We first consider the case where
$\hat{X}^{(1)},...,\hat{X}^{(K-1)}$ are rejected and see whether we accept
$\hat{X}^{(K)}$ or not; we either accept $\hat{X}^{(K)}$ with probability
$\Theta^{(K)}$ in (3) or sample a new token $Y\sim
q^{(K+1)}(\cdot|\hat{X}^{(1:K-1)})$ when all draft tokens are rejected. Since
$q^{(K+1)}$ is the residual distribution from $q^{(K)}$, one can regard it as
the simple sampling by Chen et al. (2023) and Leviathan et al. (2023), which
recovers $q^{(K)}$. The same idea is applied to
$\hat{X}^{(K-1)},...,\hat{X}^{(1)}$ in the reversed order until we recover
$q=q^{(1)}$ at the end.

Let us desribe the formal proof. From the definition of recursive rejection
sampling (4), we have

$\displaystyle\mathrm{Pr}\left\\{Z=z\right\\}=\sum_{k=1}^{K}\underbrace{\mathrm{Pr}\left\\{A^{(1:k-1)}=\mathtt{rej}^{k-1},\hat{X}^{(k)}=z,A^{(k)}=\mathtt{acc}\right\\}}_{=:\Sigma_{1,k}}+\underbrace{\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K},\tilde{X}^{(K+1)}=z\right\\}}_{=:\Sigma_{2,K}}.$

It can be shown that the following equality holds for each $k$:

$\displaystyle\Sigma_{2,k-1}=\Sigma_{1,k}+\Sigma_{2,k}.$ (13)

Let us first consider $k=K$, then,

$\displaystyle\Sigma_{1,K}+\Sigma_{2,K}$
$\displaystyle=\sum_{x^{(1)},...,x^{(K-1)}}\mathrm{Pr}\left\\{\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\times\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1},\hat{X}^{(K)}=z,A^{(K)}=\mathtt{acc}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}$
$\displaystyle+\sum_{x^{(1)},...,x^{(K)}}\mathrm{Pr}\left\\{\hat{X}^{(1:K)}=x^{(1:K)}\right\\}\times\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K},\hat{X}^{(K+1)}=z\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}$
$\displaystyle=\sum_{x^{(1)},...,x^{(K-1)}}\mathrm{Pr}\left\\{\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\Bigg{(}\underbrace{\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1},\hat{X}^{(K)}=z,A^{(K)}=\mathtt{acc}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}}_{=:T_{1}(K)}$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}+\underbrace{\sum_{x^{(K)}}\mathrm{Pr}\left\\{\hat{X}^{(K)}=x^{(K)}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\times\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K},\hat{X}^{(K+1)}=z\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}}_{=:T_{2}(K)}\Bigg{)}.$

One can represent $T_{1,K}$ and $T_{2,K}$ as follows:

$\displaystyle T_{1,K}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}$
$\displaystyle~{}~{}~{}~{}~{}\times\mathrm{Pr}\left\\{\hat{X}^{(K)}=z\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\times\mathrm{Pr}\left\\{A^{(K)}=\mathtt{acc}\Bigg{|}\hat{X}^{(1:K)}=(x^{(1:K-1)},z)\right\\}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}p^{(K)}(z|x^{(1:K-1)})\min\left\\{1,\frac{q^{(K)}(z|x^{(1:K-2)})}{p^{(K)}(z|x^{(1:K-1)})}\right\\}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\min\left\\{p^{(K)}(z|x^{(1:K-1)}),q^{(K)}(z|x^{(1:K-2)})\right\\},$
$\displaystyle T_{2,K}$
$\displaystyle=\sum_{x^{(K)}}\mathrm{Pr}\left\\{\hat{X}^{(K)}=x^{(K)}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\times\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K},\hat{X}^{(K+1)}=z\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}$
$\displaystyle=\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K}\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}\times\mathrm{Pr}\left\\{\hat{X}^{(K+1)}=z\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}$
$\displaystyle=\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\mathrm{Pr}\left\\{A^{(1:K)}=\mathtt{rej}^{K}\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}\times
q^{(K+1)}(z|x^{(1:K-1)})$
$\displaystyle=\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\mathrm{Pr}\left\\{A^{(K)}=\mathtt{rej}\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}\times
q^{(K+1)}(z|x^{(1:K-1)})$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}q^{(K+1)}(z|x^{(1:K-1)})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\mathrm{Pr}\left\\{A^{(K)}=\mathtt{rej}\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}q^{(K+1)}(z|x^{(1:K-1)})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\left(1-\mathrm{Pr}\left\\{A^{(K)}=\mathtt{acc}\bigg{|}\hat{X}^{(1:K)}=x^{(1:K)}\right\\}\right)$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}q^{(K+1)}(z|x^{(1:K-1)})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\sum_{x^{(K)}}p^{(K)}(x^{(K)}|x^{(1:K-1)})\times\left(1-\min\left\\{1,\frac{q^{(K)}(x^{(K)}|x^{(1:K-2)})}{p^{(K)}(x^{(K)}|x^{(1:K-1)})}\right\\}\right)$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\frac{\max\left\\{0,q^{(K)}(z|x^{(1:K-2)})-p^{(K)}(z|x^{(1:K-1)})\right\\}}{\sum_{x^{(K)}}\max\left\\{0,q^{(K)}(x^{(k)}|x^{(1:K-2)})-p^{(K)}(x^{(K)}|x^{(1:K-1)})\right\\}}$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\sum_{x^{(K)}}\max\left\\{0,q^{(K)}(x^{(K)}|x^{(1:K-2)})-p^{(K)}(x^{(K)}|x^{(1:K-1)})\right\\}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\max\left\\{0,q^{(K)}(z|x^{(1:K-2)})-p^{(K)}(z|x^{(1:K-1)})\right\\}.$

Therefore, we have

$\displaystyle T_{1,K}+T_{2,K}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\times\left(\min\left\\{p^{(K)}(z|x^{(1:K-1)}),q^{(K)}(z|x^{(1:K-2)})\right\\}+\max\left\\{0,q^{(K)}(z|x^{(1:K-2)})-p^{(K)}(z|x^{(1:K-1)})\right\\}\right)$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1}\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}q^{(K)}(z|x^{(1:K-2)})$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1},\tilde{X}^{(K)}=z\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\},$

where we define a random variable $\tilde{X}^{(K)}$ such that

$\displaystyle\mathrm{Pr}\left\\{\tilde{X}^{(K)}=z\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}:=q^{(K)}(z|x^{(1:K-1)}),$

which leads to

$\displaystyle\Sigma_{1,K}+\Sigma_{2,K}$
$\displaystyle=\sum_{x^{(1)},...,x^{(K-1)}}\mathrm{Pr}\left\\{\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}(T_{1,K}+T_{2,K})$
$\displaystyle=\sum_{x^{(1)},...,x^{(K-1)}}\mathrm{Pr}\left\\{\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}\times\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1},\tilde{X}^{(K)}=z\bigg{|}\hat{X}^{(1:K-1)}=x^{(1:K-1)}\right\\}$
$\displaystyle=\mathrm{Pr}\left\\{A^{(1:K-1)}=\mathtt{rej}^{K-1},\tilde{X}^{(K)}=z\right\\}$
$\displaystyle=\Sigma_{2,K-1}.$

Since the same derivation can be done for $k=2,...,K-1$, we have

$\displaystyle\mathrm{Pr}\left\\{Z=z\right\\}=\sum_{k=1}^{K}\Sigma_{1,k}+\Sigma_{2,K}=\sum_{k=1}^{K-1}\Sigma_{1,k}+\Sigma_{2,K-1}=\cdots=\Sigma_{1,1}+\Sigma_{2,1}=q(z),$

where the last equality holds from the derivation of original speculative
decoding by (Chen et al., 2023; Leviathan et al., 2023). ∎

###  A.2 Proof of Theorem 3.2

######  Theorem 3.2 (Tokens from the same sequence follow sampling without
replacement in RSD-S).

In RSD-S, any non-empty subsequence of the sequence
$\hat{X}_{l+1}^{(1)},...,\hat{X}_{l+1}^{(W)}$ of draft tokens (from
$O_{l+1}^{(1)},...,O_{l+1}^{(W)}$ in (12)) such that each element of the
subsequence has the same parent $\boldsymbol{\tau}_{l}^{(k)}$ follows sampling
without replacement from $p(\cdot|\boldsymbol{\tau}_{l}^{(k)})$.

###### Proof.

For fixed $\boldsymbol{\tau}_{l}^{(k)}$, consider a sequence of tokens

$\displaystyle\bar{\mathbf{X}}_{l+1}^{(k)}:=\underset{x\in\mathcal{X}}{\mathrm{argsort}}~{}\psi_{l}(\boldsymbol{\tau}_{l}^{(k)},x)=\underset{x\in\mathcal{X}}{\mathrm{argsort}}~{}\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},x),$

where the last equality holds since $T$ in (10) is monotonically increasing
w.r.t. $\tilde{\phi}_{l}(\boldsymbol{\tau}_{l}^{(k)},\cdot)$ for fixed
$\boldsymbol{\tau}_{l}^{(k)}$. Thus, $\bar{\mathbf{X}}_{l+1}^{(k)}$ can be
seen as samples from $p(\cdot|\boldsymbol{\tau}_{l}^{(k)})$ without
replacement.

For a length-$l_{k}$ subsequence $\mathbf{o}_{l+1}^{(k)}$ of
$(O_{l+1}^{(1)},...,O_{l+1}^{(W)})$ in (12), where each element of the
subsequence have $\boldsymbol{\tau}_{l}^{(k)}$ as its parent, the token
sequence in $\mathbf{o}_{l+1}^{(k)}$ is a subsequence of
$\bar{\mathbf{X}}_{l+1}^{(k)}$, i.e., those tokens are top-$l_{k}$ samples
without replacement from $p(\cdot|\boldsymbol{\tau}_{l}^{(k)})$. ∎

##  Appendix B Algorithm

###  B.1 Recursive Speculative Decoding with Constant Branching Factors
(RSD-C)

Algorithm 2 Recursive Speculative Decoding with Constant Branching Factors
(RSD-C)

1: Input: The length $L_{\mathrm{draft}}$ of draft sequences (depth of the
draft tree), a sequence $\mathbf{x}_{\mathrm{input}}$ of input tokens, a list
$\mathbf{b}:=[b_{0},...,b_{L_{\mathrm{draft}}-1}]$ of constant branching
factors in the draft tree, the maximum length $L_{\mathrm{output}}$ of the
output sequence.

2: // Get the length of the input sequence.
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$.

3: // Initialize empty KV caches for draft and target models.
$\mathbf{C}_{\mathrm{draft}}\leftarrow\emptyset$,
$\mathbf{C}_{\mathrm{target}}\leftarrow\emptyset$.

4: while $L_{\mathrm{input}}<L_{\mathrm{output}}$ do

5: // (STEP 1) Create a draft tree by using the draft model.
$\mathcal{T},\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{M},\mathbf{id}_{\mathrm{position}},\mathcal{L}_{\mathrm{num\\_nodes}}\leftarrow\mathtt{CreateDraftTreeConst}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{b},L_{\mathrm{draft}}).$

6: // (STEP 2) Evaluate draft tokens by using the target model. // - Apply
$\mathbf{M}$ to the right below corner of attention weights. // - The target
log probability $\boldsymbol{\Phi}_{\mathrm{target}}$ is a
$\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})\times N_{\mathrm{vocab}}$
tensor. // - $N_{\mathrm{vocab}}$ is the vocabulary size.
$\boldsymbol{\Phi}_{\mathrm{target}},\mathbf{C}_{\mathrm{target}}\leftarrow\mathtt{TargetModelForwardPass}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{target}},\mathbf{id}_{\mathrm{position}},\mathbf{M})$.

7: // - Convert the log probability tensor into the list of log probabilities
for each level of the tree.
$\mathcal{L}_{\mathrm{log\\_probs\\_target}}\leftarrow\mathtt{SplitTensor}(\boldsymbol{\Phi}_{\mathrm{target}}[-\mathtt{Sum}(\mathcal{L}_{\mathrm{num\\_nodes}}):,:],\mathcal{L}_{\mathrm{num\\_nodes}},\mathtt{dim}=0)$

8: // (STEP 3) Run Recursive Rejection Sampling for each level of the tree.
$\mathbf{x}_{\mathrm{accepted}},\mathbf{x}_{\mathrm{last}},\mathbf{id}_{\mathrm{accepted\\_flat\\_node}}\leftarrow\mathtt{RecursiveRejectionSampling}(\mathcal{T},\mathcal{L}_{\mathrm{log\\_probs\\_target}})$

9: // (STEP 4) Use KV caches that are accepted, and prepare for the next
round.
$\mathbf{C}_{\mathrm{draft}},\mathbf{C}_{\mathrm{target}}\leftarrow\mathtt{FilterKVCache}(\mathbf{C}_{\mathrm{draft}},\mathbf{C}_{\mathrm{target}},L_{\mathrm{input}},\mathbf{id}_{\mathrm{accepted\\_flat\\_node}})$

10:
$\mathbf{x}_{\mathrm{input}}\leftarrow\mathtt{Concat}([\mathbf{x}_{\mathrm{input}}[:L_{\mathrm{input}}],\mathbf{x}_{\mathrm{accepted}},\mathbf{x}_{\mathrm{last}}])$

11:
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$

12: end while

13: Output: a sequence $\mathbf{x}_{\mathrm{input}}$ that includes both input
tokens and generated output tokens.

Algorithm 3
$\mathtt{CreateDraftTreeConst}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{b},L_{\mathrm{draft}})$

1: Input: An input sequence $\mathbf{x}_{\mathrm{input}}$, the draft KV cache
$\mathbf{C}_{\mathrm{draft}}$, the branching factor
$\mathbf{b}:=[b_{0},...,b_{L_{\mathrm{draft}}-1}]$, the draft length
$L_{\mathrm{draft}}$

2: // Get the length of the input sequence.
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$.

3: // Initialize lists for 1) draft log probabilities, 2) flattened node IDs,
3) parent node ids (within each level of the draft tree), 4) draft tokens, 5)
numbers of nodes (for all levels of the tree), respectively.
$\mathcal{L}_{\mathrm{log\\_probs\\_draft}}\leftarrow[~{}],\mathcal{L}_{\mathrm{flat\\_node\\_ids}}\leftarrow[~{}],\mathcal{L}_{\mathrm{parent\\_ids}}\leftarrow[~{}],\mathcal{L}_{\mathrm{draft\\_tokens}}\leftarrow[~{}],\mathcal{L}_{\mathrm{num\\_nodes}}\leftarrow[~{}]$.

4: // Initialize a draft tree.
$\mathcal{T}\leftarrow(\mathcal{L}_{\mathrm{log\\_probs\\_draft}},\mathcal{L}_{\mathrm{flat\\_node\\_ids}},\mathcal{L}_{\mathrm{parent\\_ids}},\mathcal{L}_{\mathrm{draft\\_tokens}})$.

5: // Set an empty attention mask, and position ids; inclusive for
$\mathtt{start}$ and exclusive for $\mathtt{end}$.
$\mathbf{M}\leftarrow\emptyset$,
$\mathbf{id}_{\mathrm{position}}\leftarrow\mathtt{Arange}(\mathtt{start}=0,\mathtt{end}=L_{\mathrm{input}})$.

6: // Set the counter to check the number of nodes in the tree.
$N_{\mathrm{tree\\_prev}}\leftarrow 0$, $N_{\mathrm{tree\\_curr}}\leftarrow
0$.

7: // Set the number of nodes at the current level of the tree.
$N_{\mathrm{nodes}}\leftarrow 1$,
$\mathcal{L}_{\mathrm{num\\_nodes}}\mathtt{.append}(N_{\mathrm{nodes}})$.

8: for $l_{\mathrm{draft}}=0$ to $L_{\mathrm{draft}}-1$ do

9: // Apply $\mathbf{M}$ to the right below corner of attention weights. //
The draft log probability $\boldsymbol{\Phi}_{\mathrm{draft}}$ is a
$\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})\times N_{\mathrm{vocab}}$
tensor. // $N_{\mathrm{vocab}}$ is the vocabulary size.
$\boldsymbol{\Phi}_{\mathrm{draft}},\mathbf{C}_{\mathrm{draft}}\leftarrow\mathtt{DraftModelForwardPass}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{id}_{\mathrm{position}},\mathbf{M})$.

10: // Sample $b_{l_{\mathrm{draft}}}$ nodes without replacement,
independently for $N_{\mathrm{nodes}}$ nodes.// NOTE: Outputs are sorted
w.r.t. the value of perturbed log probabilities and
flattened.$\mathbf{x}_{\mathrm{draft}},\mathbf{id}_{\mathrm{parent}}\leftarrow\mathtt{SampleWithGumbelTopK}(\boldsymbol{\Phi}_{\mathrm{draft}}[-N_{\mathrm{nodes}}:,:],b_{l_{\mathrm{draft}}})$.

11: // Update the input sequence of
tokens.$\mathbf{x}_{\mathrm{input}}\leftarrow\mathtt{Concat}([\mathbf{x}_{\mathrm{input}},\mathbf{x}_{\mathrm{draft}}])$.

12: // Get the number of newly added
nodes.$N_{\mathrm{nodes}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{draft}})$.

13: // Build attention mask reflecting tree
topology.$\mathbf{M}\leftarrow\mathtt{BuildAttentionMask}(\mathbf{M},\mathbf{id}_{\mathrm{parent}},N_{\mathrm{nodes}},N_{\mathrm{tree\\_prev}},N_{\mathrm{tree\\_curr}})$.

14: // Update counters.$N_{\mathrm{tree\\_prev}}\leftarrow
N_{\mathrm{tree\\_curr}}$, $N_{\mathrm{tree\\_curr}}\leftarrow
N_{\mathrm{tree\\_curr}}+N_{\mathrm{nodes}}$.

15: // Update position IDs.
$\mathbf{id}_{\mathrm{position}}\leftarrow\mathtt{Concat}([\mathbf{id}_{\mathrm{position}},(L_{\mathrm{input}}+l_{\mathrm{draft}})\times\mathbf{1}_{N_{\mathrm{nodes}}}])$.

16: // Get node IDs considering the flattened draft tree.// This is used to
update KV
caches.$\mathbf{id}_{\mathrm{flat\\_node}}\leftarrow\mathtt{Arange}(\mathtt{start}=L_{\mathrm{input}}+N_{\mathrm{tree\\_prev}},\mathtt{end}=L_{\mathrm{input}}+N_{\mathrm{tree\\_curr}})$.

17: // Update the lists of the
tree.$\mathcal{L}_{\mathrm{log\\_probs\\_draft}}\mathtt{.append}(\boldsymbol{\Phi}_{\mathrm{draft}})$,
$\mathcal{L}_{\mathrm{flat\\_node\\_ids}}\mathtt{.append}(\mathbf{id}_{\mathrm{flat\\_node}})$,
$\mathcal{L}_{\mathrm{parent\\_ids}}\mathtt{.append}(\mathbf{id}_{\mathrm{parent}})$,$\mathcal{L}_{\mathrm{draft\\_tokens}}\mathtt{.append}(\mathbf{x}_{\mathrm{draft}})$,
$\mathcal{L}_{\mathrm{num\\_nodes}}\mathtt{.append}(N_{\mathrm{nodes}})$.

18: end for

19: Output:
$\mathcal{T},\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{M},\mathbf{id}_{\mathrm{position}},\mathcal{L}_{\mathrm{num\\_nodes}}$.

Algorithm 4 $\mathtt{SampleWithGumbelTopK}(\boldsymbol{\Phi},K)$

1: Input: a $N_{\mathrm{nodes}}\times N_{\mathrm{vocab}}$ log probabilities
$\boldsymbol{\Phi}$, the number $K$ of desired samples without replacement.

2: // Sample a matrix where elements are i.i.d. standard Gumbel random
variables.$\mathbf{G}\leftarrow[g_{ij}]$,
$g_{ij}\leftarrow\mathtt{SampleStandardGumbel}(),i=0,...,N_{\mathrm{nodes}}-1,j=0,...,N_{\mathrm{vocab}}-1$.

3: // Perturb log probabilities with Gumbel random
variables.$\tilde{\boldsymbol{\Phi}}\leftarrow\boldsymbol{\Phi}+\mathbf{G}$.

4: // Get top-$K$ elements corresponding to the $K$ largest perturb log
probabilities.// Outputs are sorted (in each row) w.r.t. the values of
perturbed log probabilities and
flattened.$\mathbf{x}\leftarrow\mathrm{argtop}^{(K)}(\tilde{\boldsymbol{\Phi}},\mathtt{dim}=-1)\mathtt{.flatten}()$.

5: // Set parent
ids.$\mathbf{id}_{\mathrm{parent}}\leftarrow\mathtt{Concat}([0\cdot\mathbf{1}_{K},1\cdot\mathbf{1}_{K},...,(N_{\mathrm{nodes}}-1)\cdot\mathbf{1}_{K}])$.

6: // When probability filtering methods(e.g., top-$p$, top-$k$) were applied,
filter some elements of $\mathbf{x}$ and $\mathbf{id}_{\mathrm{parent}}$ if
corresponding log probability is equal to $-\infty$.

7: Output: $\mathbf{x},\mathbf{id}_{\mathrm{parent}}$.

Algorithm 5
$\mathtt{BuildAttentionMask}(\mathbf{M},\mathbf{id}_{\mathrm{parent}},N_{\mathrm{nodes}},N_{\mathrm{tree\\_prev}},N_{\mathrm{tree\\_curr}})$

1: Input: previous attention mask $\mathbf{M}$, parent node ids
$\mathbf{id}_{\mathrm{parent}}$ for newly added nodes, the number
$N_{\mathrm{nodes}}$ of nodes newly added to the tree, the total number
$N_{\mathrm{tree\\_prev}}$ of nodes in the previous-iteration tree, the total
number $N_{\mathrm{tree\\_curr}}$ of nodes in the current-iteration tree.

2: if $\mathbf{M}==\emptyset$ then

3: // If the attention mask is empty, we initialize with zeros.
$\mathbf{M}\leftarrow\mathbf{0}_{N_{\mathrm{nodes}}\times
N_{\mathrm{nodes}}}$.

4: else

5: // If the attention mask exists, we zero-pad.
$\mathbf{M}\leftarrow\mathtt{ZeroPadding}(\mathbf{M},\mathtt{right}=N_{\mathrm{nodes}},\mathtt{bottom}=N_{\mathrm{nodes}})$.

6: for $i=0$ to $N_{\mathrm{nodes}}-1$ do

7: // Copy the row about paraent nodes to the row about child nodes.
$\mathbf{M}[N_{\mathrm{tree\\_curr}}+i,:]\leftarrow\mathbf{M}[N_{\mathrm{tree\\_prev}}+\mathbf{id}_{\mathrm{parent}}[i],:]$.

8: end for

9: end if

10: // Set diagonal elements equal to 1.
$\mathbf{M}\leftarrow\mathbf{M}\mathtt{.fill\\_diagonal}(1)$

11: Output: the new attention mask $\mathbf{M}$.

Algorithm 6
$\mathtt{RecursiveRejectionSampling}(\mathcal{T},\mathcal{L}_{\mathrm{log\\_probs\\_target}})$

1: Input: the draft tree $\mathcal{T}$, the list
$\mathcal{L}_{\mathrm{log\\_probs\\_target}}$ of target log probabilities

2: // Get lists from the draft tree.
$\mathcal{L}_{\mathrm{log\\_probs\\_draft}},\mathcal{L}_{\mathrm{flat\\_node\\_ids}},\mathcal{L}_{\mathrm{parent\\_ids}},\mathcal{L}_{\mathrm{draft\\_tokens}}\leftarrow\mathcal{T}$

3: // Set the current node id. $i_{node}\leftarrow 0$

4: // Initialize the lists to store accepted draft tokens and flattened node
ids (for KV cache update).
$\mathcal{L}_{\mathrm{accepted\\_draft\\_tokens}}\leftarrow[~{}],\mathcal{L}_{\mathrm{accepted\\_flat\\_node\\_ids}}\leftarrow[~{}]$.

5: for $l_{\mathrm{draft}}=0$ to $L_{\mathrm{draft}}-1$ do

6: // Get log probabilities at the current node. // Both are $1\times
N_{\mathrm{vocab}}$ tensors, where $N_{\mathrm{vocab}}$ is the vocabulary
size.
$\boldsymbol{\Phi}_{\mathrm{draft}}\leftarrow\mathcal{L}_{\mathrm{log\\_probs\\_draft}}[l_{\mathrm{draft}}][i_{node}:(i_{node}+1),:]$,
$\boldsymbol{\Phi}_{\mathrm{target}}\leftarrow\mathcal{L}_{\mathrm{log\\_probs\\_target}}[l_{\mathrm{draft}}][i_{node}:(i_{node}+1),:]$

7: // Get draft tokens, flattened node IDs, parent IDs at the current level.
$\mathbf{x}_{\mathrm{draft}}\leftarrow\mathcal{L}_{\mathrm{draft\\_tokens}}[l_{\mathrm{draft}}]$,
$\mathbf{id}_{\mathrm{flat\\_node}}\leftarrow\mathcal{L}_{\mathrm{flat\\_node\\_ids}}[l_{\mathrm{draft}}]$,
$\mathbf{id}_{\mathrm{parent}}\leftarrow\mathcal{L}_{\mathrm{parent\\_ids}}[l_{\mathrm{draft}}]$

8: // Initialize an acceptance indicator as False.
$\mathrm{accept}\leftarrow\mathrm{False}$

9: for $i$ in $\mathbf{id}_{\mathrm{parent}}$ do

10: if $i\neq i_{node}$ then

11: continue

12: end if

13: // Get the current draft token.
$x_{d}\leftarrow\mathbf{x}_{\mathrm{draft}}[i]$.

14: // Sample a uniform random variable. $U\sim\mathrm{Uniform}[0,1]$.

15: if
$U<\min\\{1,\exp(\boldsymbol{\Phi}_{\mathrm{target}}[0,x_{d}]-\boldsymbol{\Phi}_{\mathrm{draft}}[0,x_{d}])\\}$
then

16: // Set the indicator as True is the token is accepted.
$\mathrm{accept}\leftarrow\mathrm{True}$.

17: // Store the accepted token and corresponding flattened node ID.
$\mathcal{L}_{\mathrm{accepted\\_draft\\_tokens}}\mathtt{.append}(x_{d})$,
$\mathcal{L}_{\mathrm{accepted\\_draft\\_tokens}}\mathtt{.append}(\mathbf{id}_{\mathrm{flat\\_node}}[i])$.

18: $i_{node}\leftarrow i$.

19: break

20: end if

21: // Get clamped target log probability.
$\boldsymbol{\Phi}_{\mathrm{target}}\leftarrow\log((\exp(\boldsymbol{\Phi}_{\mathrm{target}})-\exp(\boldsymbol{\Phi}_{\mathrm{draft}}))\mathtt{.clamp}(\mathtt{min}=0)\\})$

22: // Normalize the clamped target log probability.
$\boldsymbol{\Phi}_{\mathrm{target}}\leftarrow\boldsymbol{\Phi}_{\mathrm{target}}-\mathtt{LogSumExp}(\boldsymbol{\Phi}_{\mathrm{target}})$

23: // Neglect draft log probability of already sampled token.
$\boldsymbol{\Phi}_{\mathrm{draft}}[i]\leftarrow-\infty$

24: // Normalize the draft log probability.
$\boldsymbol{\Phi}_{\mathrm{draft}}\leftarrow\boldsymbol{\Phi}_{\mathrm{draft}}-\mathtt{LogSumExp}(\boldsymbol{\Phi}_{\mathrm{draft}})$

25: end for

26: if $\mathrm{accept}==\mathrm{False}$ then

27: break

28: end if

29: end for

30: if $\mathrm{accept}$ then

31: // At the leaf node when all tokens are accepted, we use target log
probability to draw a sample.
$\boldsymbol{\Phi}_{\mathrm{target}}\leftarrow\mathcal{L}_{\mathrm{log\\_probs\\_target}}[l_{d}][i_{node}:(i_{node}+1),:]$

32: end if

33:
$\mathbf{x}_{\mathrm{last}}\sim\mathtt{SampleWithGumbelTopK}(\boldsymbol{\Phi}_{\mathrm{target}},1)$

34:
$\mathbf{x}_{\mathrm{accepted}}\leftarrow\mathtt{Stack}(\mathcal{L}_{\mathrm{accepted\\_draft\\_tokens}})$,
$\mathbf{id}_{\mathrm{accepted\\_flat\\_node}}\leftarrow\mathtt{Stack}(\mathcal{L}_{\mathrm{accepted\\_draft\\_tokens}})$.

35: Output:
$\mathbf{x}_{\mathrm{accepted}},\mathbf{x}_{\mathrm{last}},\mathbf{id}_{\mathrm{accepted\\_flat\\_node}}$

###  B.2 Recursive Speculative Decoding with Stochastic Beam Search (RSD-S)

We highlight the difference w.r.t. RSD-C.

Algorithm 7 Recursive Speculative Decoding with Stochastic Beam Search (RSD-S)

1: Input: The length $L_{\mathrm{draft}}$ of draft sequences (depth of the
draft tree), a sequence $\mathbf{x}_{\mathrm{input}}$ of input tokens, the
beamwidth $W$, the maximum length $L_{\mathrm{output}}$ of the output
sequence.

2: // Get the length of the input sequence.
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$.

3: // Initialize empty KV caches for draft and target models.
$\mathbf{C}_{\mathrm{draft}}\leftarrow\emptyset$,
$\mathbf{C}_{\mathrm{target}}\leftarrow\emptyset$.

4: while $L_{\mathrm{input}}<L_{\mathrm{output}}$ do

5: // (STEP 1) Create a draft tree by using the draft model.
$\mathcal{T},\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{M},\mathbf{id}_{\mathrm{position}},\mathcal{L}_{\mathrm{num\\_nodes}}$
$\leftarrow{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathtt{CreateDraftTreeStochasticBeamSearch}}}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}W},L_{\mathrm{draft}}).$

6: // (STEP 2) Evaluate draft tokens by using the target model. // - Apply
$\mathbf{M}$ to the right below corner of attention weights. // - The target
log probability $\boldsymbol{\Phi}_{\mathrm{target}}$ is a
$\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})\times N_{\mathrm{vocab}}$
tensor. // - $N_{\mathrm{vocab}}$ is the vocabulary size.
$\boldsymbol{\Phi}_{\mathrm{target}},\mathbf{C}_{\mathrm{target}}\leftarrow\mathtt{TargetModelForwardPass}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{target}},\mathbf{id}_{\mathrm{position}},\mathbf{M})$.

7: // - Convert the log probability tensor into the list of log probabilities
for each level of the tree.
$\mathcal{L}_{\mathrm{log\\_probs\\_target}}\leftarrow\mathtt{SplitTensor}(\boldsymbol{\Phi}_{\mathrm{target}}[-\mathtt{Sum}(\mathcal{L}_{\mathrm{num\\_nodes}}):,:],\mathcal{L}_{\mathrm{num\\_nodes}},\mathtt{dim}=0)$

8: // (STEP 3) Run Recursive Rejection Sampling for each level of the tree.
$\mathbf{x}_{\mathrm{accepted}},\mathbf{x}_{\mathrm{last}},\mathbf{id}_{\mathrm{accepted\\_flat\\_node}}\leftarrow\mathtt{RecursiveRejectionSampling}(\mathcal{T},\mathcal{L}_{\mathrm{log\\_probs\\_target}})$

9: // (STEP 4) Use KV caches that are accepted, and prepare for the next
round.
$\mathbf{C}_{\mathrm{draft}},\mathbf{C}_{\mathrm{target}}\leftarrow\mathtt{FilterKVCache}(\mathbf{C}_{\mathrm{draft}},\mathbf{C}_{\mathrm{target}},L_{\mathrm{input}},\mathbf{id}_{\mathrm{accepted\\_flat\\_node}})$

10:
$\mathbf{x}_{\mathrm{input}}\leftarrow\mathtt{Concat}([\mathbf{x}_{\mathrm{input}}[:L_{\mathrm{input}}],\mathbf{x}_{\mathrm{accepted}},\mathbf{x}_{\mathrm{last}}])$

11:
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$

12: end while

13: Output: a sequence $\mathbf{x}_{\mathrm{input}}$ that includes both input
tokens and generated output tokens.

Algorithm 8
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathtt{CreateDraftTreeStochasticBeamSearch}}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}W},L_{\mathrm{draft}})$

1: Input: An input sequence $\mathbf{x}_{\mathrm{input}}$, the draft KV cache
$\mathbf{C}_{\mathrm{draft}}$, the beamwidth $W$, the draft length
$L_{\mathrm{draft}}$

2: // Get the length of the input sequence.
$L_{\mathrm{input}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})$.

3: // Initialize lists for 1) draft log probabilities, 2) flattened node IDs,
3) parent node ids (within each level of the draft tree), 4) draft tokens, 5)
numbers of nodes (for all levels of the tree), respectively.
$\mathcal{L}_{\mathrm{log\\_probs\\_draft}}\leftarrow[~{}],\mathcal{L}_{\mathrm{flat\\_node\\_ids}}\leftarrow[~{}],\mathcal{L}_{\mathrm{parent\\_ids}}\leftarrow[~{}],\mathcal{L}_{\mathrm{draft\\_tokens}}\leftarrow[~{}],\mathcal{L}_{\mathrm{num\\_nodes}}\leftarrow[~{}]$.

4: // Initialize a draft tree.
$\mathcal{T}\leftarrow(\mathcal{L}_{\mathrm{log\\_probs\\_draft}},\mathcal{L}_{\mathrm{flat\\_node\\_ids}},\mathcal{L}_{\mathrm{parent\\_ids}},\mathcal{L}_{\mathrm{draft\\_tokens}})$.

5: // Set an empty attention mask, and position ids; inclusive for
$\mathtt{start}$ and exclusive for $\mathtt{end}$.
$\mathbf{M}\leftarrow\emptyset$,
$\mathbf{id}_{\mathrm{position}}\leftarrow\mathtt{Arange}(\mathtt{start}=0,\mathtt{end}=L_{\mathrm{input}})$.

6: // Set the counter to check the number of nodes in the tree.
$N_{\mathrm{tree\\_prev}}\leftarrow 0$, $N_{\mathrm{tree\\_curr}}\leftarrow
0$.

7: // Set the number of nodes at the current level of the tree.
$N_{\mathrm{nodes}}\leftarrow 1$,
$\mathcal{L}_{\mathrm{num\\_nodes}}\mathtt{.append}(N_{\mathrm{nodes}})$.

8: // Set stochastic beam parameters: sum log probabilities
$\boldsymbol{\Sigma}$ and truncated Gumbels $\boldsymbol{\Gamma}$ for each
node in the current level of draft tree
$\boldsymbol{\Sigma}\leftarrow\textbf{0}_{N_{\mathrm{nodes}}\times
1},\boldsymbol{\Gamma}\leftarrow\textbf{0}_{N_{\mathrm{nodes}}\times 1}$.

9: for $l_{\mathrm{draft}}=0$ to $L_{\mathrm{draft}}-1$ do

10: // Apply $\mathbf{M}$ to the right below corner of attention weights. //
The draft log probability $\boldsymbol{\Phi}_{\mathrm{draft}}$ is a
$\mathtt{GetLength}(\mathbf{x}_{\mathrm{input}})\times N_{\mathrm{vocab}}$
tensor. // $N_{\mathrm{vocab}}$ is the vocabulary size.
$\boldsymbol{\Phi}_{\mathrm{draft}},\mathbf{C}_{\mathrm{draft}}\leftarrow\mathtt{DraftModelForwardPass}(\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{id}_{\mathrm{position}},\mathbf{M})$.

11: // Sample $b_{l_{\mathrm{draft}}}$ nodes without replacement,
independently for $N_{\mathrm{nodes}}$ nodes.// NOTE: Outputs are sorted
w.r.t. the value of perturbed log probabilities and
flattened.$\mathbf{x}_{\mathrm{draft}},\mathbf{id}_{\mathrm{parent}},\boldsymbol{\Sigma},\boldsymbol{\Gamma}\leftarrow\mathtt{SampleWithStochasticBeam}(\boldsymbol{\Phi}_{\mathrm{draft}}[-N_{\mathrm{nodes}}:,:],\boldsymbol{\Sigma},\boldsymbol{\Gamma},W)$.

12: // Update the input sequence of
tokens.$\mathbf{x}_{\mathrm{input}}\leftarrow\mathtt{Concat}([\mathbf{x}_{\mathrm{input}},\mathbf{x}_{\mathrm{draft}}])$.

13: // Get the number of newly added
nodes.$N_{\mathrm{nodes}}\leftarrow\mathtt{GetLength}(\mathbf{x}_{\mathrm{draft}})$.

14: // Build attention mask reflecting tree
topology.$\mathbf{M}\leftarrow\mathtt{BuildAttentionMask}(\mathbf{M},\mathbf{id}_{\mathrm{parent}},N_{\mathrm{nodes}},N_{\mathrm{tree\\_prev}},N_{\mathrm{tree\\_curr}})$.

15: // Update counters.$N_{\mathrm{tree\\_prev}}\leftarrow
N_{\mathrm{tree\\_curr}}$, $N_{\mathrm{tree\\_curr}}\leftarrow
N_{\mathrm{tree\\_curr}}+N_{\mathrm{nodes}}$.

16: // Update position IDs.
$\mathbf{id}_{\mathrm{position}}\leftarrow\mathtt{Concat}([\mathbf{id}_{\mathrm{position}},(L_{\mathrm{input}}+l_{\mathrm{draft}})\times\mathbf{1}_{N_{\mathrm{nodes}}}])$.

17: // Get node IDs considering the flattened draft tree.// This is used to
update KV
caches.$\mathbf{id}_{\mathrm{flat\\_node}}\leftarrow\mathtt{Arange}(\mathtt{start}=L_{\mathrm{input}}+N_{\mathrm{tree\\_prev}},\mathtt{end}=L_{\mathrm{input}}+N_{\mathrm{tree\\_curr}})$.

18: // Update the lists of the
tree.$\mathcal{L}_{\mathrm{log\\_probs\\_draft}}\mathtt{.append}(\boldsymbol{\Phi}_{\mathrm{draft}})$,
$\mathcal{L}_{\mathrm{flat\\_node\\_ids}}\mathtt{.append}(\mathbf{id}_{\mathrm{flat\\_node}})$,
$\mathcal{L}_{\mathrm{parent\\_ids}}\mathtt{.append}(\mathbf{id}_{\mathrm{parent}})$,$\mathcal{L}_{\mathrm{draft\\_tokens}}\mathtt{.append}(\mathbf{x}_{\mathrm{draft}})$,
$\mathcal{L}_{\mathrm{num\\_nodes}}\mathtt{.append}(N_{\mathrm{nodes}})$.

19: end for

20: Output:
$\mathcal{T},\mathbf{x}_{\mathrm{input}},\mathbf{C}_{\mathrm{draft}},\mathbf{M},\mathbf{id}_{\mathrm{position}},\mathcal{L}_{\mathrm{num\\_nodes}}$.

Algorithm 9
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\mathtt{SampleWithStochasticBeam}(\boldsymbol{\Phi},\boldsymbol{\Sigma},\boldsymbol{\Gamma},K)}$

1: Input: a $N_{\mathrm{nodes}}\times N_{\mathrm{vocab}}$ log probabilities
$\boldsymbol{\Phi}$, a $N_{\mathrm{nodes}}\times 1$ sum log probabilities
$\boldsymbol{\Sigma}$, a $N_{\mathrm{nodes}}\times 1$ truncated Gumbels
$\boldsymbol{\Gamma}$, the beamwidth $K$.

2: // Get sum log probs up to child
nodes.$\boldsymbol{\Phi}\leftarrow\boldsymbol{\Phi}+\boldsymbol{\Sigma}\mathbf{1}_{1\times
N_{\mathrm{vocab}}}$.

3: // Sample a matrix where elements are i.i.d. standard Gumbel random
variables.$\mathbf{G}\leftarrow[g_{ij}]$,
$g_{ij}\leftarrow\mathtt{SampleStandardGumbel}(),i=0,...,N_{\mathrm{nodes}}-1,j=0,...,N_{\mathrm{vocab}}-1$.

4: // Perturb sum log probabilities with Gumbel random
variables.$\tilde{\boldsymbol{\Phi}}\leftarrow\boldsymbol{\Phi}+\mathbf{G}$.

5: // Compute row-wise maximum value of perturbed sum log probabilities.// The
output size is $N_{\mathrm{nodes}}\times
1$.$\tilde{\boldsymbol{\Phi}}_{\mathrm{max}}\leftarrow\tilde{\boldsymbol{\Phi}}.\mathtt{max}(\mathtt{dim}=-1,\mathtt{keepdim}=\mathrm{True})$.

6: // Get truncated Gumbels for all expansion.// The output size is
$N_{\mathrm{nodes}}\times N_{\mathrm{vocab}}$.// NOTE: the numerical stable
way of computing this quantity was described in the original Stochastic Beam
Search
paper.$\tilde{\boldsymbol{\Gamma}}\leftarrow-\log(\exp(-\boldsymbol{\Gamma}\mathbf{1}_{1\times
N_{\mathrm{vocab}}})-\exp(-\tilde{\boldsymbol{\Phi}}_{\mathrm{max}}\mathbf{1}_{1\times
N_{\mathrm{vocab}}})+\exp(-\tilde{\boldsymbol{\Phi}}))$

7: // Get top-$K$ elements and the $K$ largest truncated Gumbels.// NOTE: we
consider top-$K$ elements for all elements in $\tilde{\boldsymbol{\Gamma}}$,
so both parent node IDs and token IDs can be acquired. Make sure that both
output IDs are sorted w.r.t. the corresponding values in
$\tilde{\boldsymbol{\Gamma}}$.$\mathbf{id}_{\mathrm{parent}},\mathbf{x},\boldsymbol{\Gamma}\leftarrow\mathrm{argtop}\text{-}K(\tilde{\boldsymbol{\Phi}})$.

8: // Get sum log probs for top-$K$
elements.$\boldsymbol{\Sigma}\leftarrow\boldsymbol{\Phi}[\mathbf{id}_{\mathrm{parent}},\mathbf{x}]$.

9: // When probability filtering methods(e.g., top-$p$, top-$k$) were applied,
filter some elements of
$\mathbf{x},\mathbf{id}_{\mathrm{parent}},\boldsymbol{\Sigma},\boldsymbol{\Gamma}$
if corresponding log probability is equal to $-\infty$.

10: Output:
$\mathbf{x},\mathbf{id}_{\mathrm{parent}},\boldsymbol{\Sigma},\boldsymbol{\Gamma}$

##  Appendix C Experiments

###  C.1 Draft models

The following draft models are used:

  * •

For Llama 2 target models, we use the 115M Llama 2 drafter and Llama 2-Chat
drafter for Llama 2 and Llama 2-Chat target models, respectively.

    * –

Llama 2 drafter uses smaller Llama archiecture (Touvron et al., 2023) and is
pre-trained on the 600B-token dataset

    * –

Llama 2-Chat drafter is the model fine-tuned from Llama 2-drafter so that it
can be aligned with Llama 2-Chat-7B via distillation.

  * •

For OPT target models, we use OPT with 125M and 350M parameters for target OPT
models.

###  C.2 Performance Metrics

In the experiments, we consider three metrics (except accuracy) for all target
models.

  * •

Block efficiency (Leviathan et al., 2023) is the average number of tokens
generated per target model call. Within a single target call, auto-regressive
decoding always generates a single token, while speculative decoding methods
generates

$\displaystyle\text{(Number of accepted tokens)}+1.$

The block efficiency $\eta$ is the average over all target calls.

  * •

Memory-Bound Speed Up (MBSU) is the fictitious inference speed-up relative to
auto-regressive decoding, where we assume each model’s runtime is proportional
to the model size. Specifically, let $L$ denote the (maximum) length of draft
sequences, which is the depth of the draft-token tree for tree-based
speculative decoding methods, and $r$ denote the relative speed of running the
draft model to that of the target model. The walltime improvement (Leviathan
et al., 2023; Zhou et al., 2023) is

$\displaystyle\frac{\eta}{L\times r+1}.$

MBSU considers a specific case where $r$ is equal to $\text{(Size of the
target model)}/\text{(Size of the draft model)}$, considering practical
scenarios in memory-bound devices where loading model weights takes
significant amount time, often proportional to their size.

  * •

Token rate is the measure of average number of generated tokens per second
while running on A100 GPUs. It shows different results from MBSU since running
A100 GPUs is far from memory-bound scenarios.

###  C.3 Tree Structure

####  C.3.1 Experiment for various lengths of draft sequence

The following trees are used for draft sequence length $L$, where SD uses a
single draft sequence with length $L$. For each $L$, we first set RSD-C with
constant branching factors always equal to 2 and set the draft-tree sizes for
SpecTr and RSD-S _always less than or equal to_ the tree size of RSD-C. Then,
we add RSD-C with the branching factor $\mathbf{b}:=[n,1,...,1]$ where $n$ is
properly set to have the draft-tree size equal to that of SpecTr and RSD-S. In
_Figure 4_, we show the best results across all tree structures for each $L$
and algorithm.

  * •

$L=2$:

    * –

SpecTr and RSD-S: $(K,L)\in\\{(2,2),(3,2)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,2],[2,1],[3,1]\\}$ for a vector $\mathbf{b}$ of
branching factors.

  * •

$L=3$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(3,3),(4,3)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,2,2],[3,1,1],[4,1,1]\\}$ for a vector $\mathbf{b}$
of branching factors.

  * •

$L=4$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(5,4),(7,4)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,2,2,2],[5,1,1,1],[7,1,1,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

  * •

$L=5$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(6,5),(12,5)\\}$, where $K$ becomes the number
of independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,2,2,2,2],[6,1,1,1,1],[12,1,1,1,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

####  C.3.2 Experiment for vairous target computational budget

The following trees are used for target computational budgets $B$, i.e., the
number of tokens to process at the target model, where $B$ becomes the draft
length of SD. In _Figure 5_, we show the best results across all tree
structures for each $B$ and algorithm.

  * •

$B=6$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(2,3),(3,2)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,1,1],[2,2],[3,1]\\}$ for a vector $\mathbf{b}$ of
branching factors.

  * •

$B=10$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(2,5),(5,2)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,1,1,1,1],[2,2,1],[5,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

  * •

$B=14$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(2,7),(7,2)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,1,1,1,1,1,1],[2,2,2],[7,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

  * •

$B=21$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(3,7),(7,3)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[3,1,1,1,1,1,1],[3,2,2],[7,1,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

  * •

$B=30$

    * –

SpecTr and RSD-S: $(K,L)\in\\{(5,6),(6,5)\\}$, where $K$ becomes the number of
independent draft sequences for SpecTr and the beamwidth for RSD-S

    * –

RSD-C: $\mathbf{b}\in\\{[2,2,2,2],[5,1,1,1,1,1],[6,1,1,1,1]\\}$ for a vector
$\mathbf{b}$ of branching factors.

###  C.4 Experiment results with plots

####  C.4.1 Block efficiency, MBSU, token rate and accuracy for various
lengths of draft sequence

Llama 2-7BLlama 2-13BLlama 2-Chat-7BLlama 2-Chat-13B Figure 6: Block
efficiency, MBSU, token rate and accuracy for varying lengths of draft
sequence are given for multiple target models: Llama 2-7B, Llama 2-13B, Llama
2-Chat-7B, Llama 2-Chat-13B. Chat models use the same draft model, while the
other models use the same draft model different from the one for chat models.
All results are normalized w.r.t. the values of AR decoding.
OPT-125M-13BOPT-125M-30BOPT-125M-66BOPT-350M-13BOPT-350M-30BOPT-350M-66B
Figure 7: Block efficiency, MBSU, token rate and accuracy for varying lengths
of draft sequence are given for multiple pairs of draft and target models: the
size of draft model is in $\\{\text{125M, 350M}\\}$, and the size of target
model is in $\\{\text{13B, 30B, 66B}\\}$. All results are normalized w.r.t.
the values of AR decoding.

####  C.4.2 Block efficiency, MBSU, token rate and accuracy for various target
computational budget

Llama 2-7BLlama 2-13BLlama 2-Chat-7BLlama 2-Chat-13B Figure 8: Block
efficiency, MBSU, token rate and accuracy for varying numbers of tokens
processed at the target model are given for multiple target models: Llama
2-7B, Llama 2-13B, Llama 2-Chat-7B, Llama 2-Chat-13B. Chat models use the same
draft model, while the other models use the same draft model different from
the one for chat models. All results are normalized w.r.t. the values of AR
decoding.
OPT-125M-13BOPT-125M-30BOPT-125M-66BOPT-350M-13BOPT-350M-30BOPT-350M-66B
Figure 9: Block efficiency, MBSU, token rate and accuracy for varying numbers
of tokens processed at the target model are given for multiple pairs of draft
and target models: the size of draft model is in $\\{\text{125M, 350M}\\}$,
and the size of target model is in $\\{\text{13B, 30B, 66B}\\}$. All results
are normalized w.r.t. the values of AR decoding.

###  C.5 Experient results with tables

For readers curious about raw numbers, we remain all the numbers used for
plots as tables in this section.

####  C.5.1 Block efficiency, MBSU, token rate and accuracy for varying
lengths of draft sequence

  * •

Llama 2-7B (with 115M drafter)

    * –

XSum (_Table_ 1), WMT (_Table_ 2)

  * •

Llama 2-13B (with 115M drafter)

    * –

XSum (_Table_ 3), WMT (_Table_ 4)

  * •

Llama 2-70B (with 115M drafter)

    * –

XSum (_Table_ 5), WMT (_Table_ 6)

  * •

Llama 2-Chat-7B (with 115M drafter)

    * –

XSum (_Table_ 7), WMT (_Table_ 8), Dolly (_Table_ 9)

  * •

Llama 2-Chat-13B (with 115M drafter)

    * –

XSum (_Table_ 10), WMT (_Table_ 11), Dolly (_Table_ 12)

  * •

Llama 2-Chat-70B (with 115M drafter)

    * –

XSum (_Table_ 13), WMT (_Table_ 14), Dolly (_Table_ 15)

  * •

OPT-13B (with OPT-125M drafter)

    * –

XSum (_Table_ 16), WMT (_Table_ 17)

  * •

OPT-30B (with OPT-125M drafter)

    * –

XSum (_Table_ 18), WMT (_Table_ 19)

  * •

OPT-66B (with OPT-125M drafter)

    * –

XSum (_Table_ 20), WMT (_Table_ 21)

  * •

OPT-13B (with OPT-350M drafter)

    * –

XSum (_Table_ 22), WMT (_Table_ 23)

  * •

OPT-30B (with OPT-350M drafter)

    * –

XSum (_Table_ 24), WMT (_Table_ 25)

  * •

OPT-66B (with OPT-350M drafter)

    * –

XSum (_Table_ 26), WMT (_Table_ 27)

####  C.5.2 Block efficiency, MBSU, token rate and accuracy for varying
numbers of tokens processed at the target model

  * •

Llama 2-7B (with 115M drafter)

    * –

XSum (_Table_ 28), WMT (_Table_ 29)

  * •

Llama 2-13B (with 115M drafter)

    * –

XSum (_Table_ 30), WMT (_Table_ 31)

  * •

Llama 2-70B (with 115M drafter)

    * –

XSum (_Table_ 32), WMT (_Table_ 33)

  * •

Llama 2-Chat-7B (with 115M drafter)

    * –

XSum (_Table_ 34), WMT (_Table_ 35), Dolly (_Table_ 36)

  * •

Llama 2-Chat-13B (with 115M drafter)

    * –

XSum (_Table_ 37), WMT (_Table_ 38), Dolly (_Table_ 39)

  * •

Llama 2-Chat-70B (with 115M drafter)

    * –

XSum (_Table_ 40), WMT (_Table_ 41), Dolly (_Table_ 42)

  * •

OPT-13B (with OPT-125M drafter)

    * –

XSum (_Table_ 43), WMT (_Table_ 44)

  * •

OPT-30B (with OPT-125M drafter)

    * –

XSum (_Table_ 45), WMT (_Table_ 46)

  * •

OPT-66B (with OPT-125M drafter)

    * –

XSum (_Table_ 47), WMT (_Table_ 48)

  * •

OPT-13B (with OPT-350M drafter)

    * –

XSum (_Table_ 49), WMT (_Table_ 50)

  * •

OPT-30B (with OPT-350M drafter)

    * –

XSum (_Table_ 51), WMT (_Table_ 52)

  * •

OPT-66B (with OPT-350M drafter)

    * –

XSum (_Table_ 53), WMT (_Table_ 54)

Table 1: We summarize experiment results with Llama 2-7B target and 115M draft
for the XSum task with various draft lengths. Draft Length (DL) means the
(maximum) length for all draft token sequences generated by the draft model.
For decoders (Dec.), we consider Auto-Regressive sampling (AR), Speculative
Decoding (SD), SpecTr, Recursive Speculative Decoding with Constant branching
factors (RSD-C), Recursive Speculative Decoding with Stochastic Beam Search
(RSD-S). The contents in decoder specification (Spec.) have different meanings
for each decoder; $K\times L$ means the number $K$ of draft paths and draft
length $L$ for SpecTr; it describes constant branching factors for each level
of the tree (from root to leaf) for RSD-C; $K\times L$ means the beamwidth $K$
and draft length $L$ for RSD-S. Block efficiency (Eff.), Memory Bound Speed Up
(MBSU), Token Rate (TR), and Accuracy (Acc.) are given for each algorithm. For
each group of rows having the same complexity, we highlight the best values
for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-7B | XSum | 0 | AR | - | 1.000 | 1.000 | 37.269 | 0.141
2 | SD | 2 | 2.166 | 2.093 | 51.470 | 0.143
SpecTr | 2$\times$2 | 2.218 | 2.143 | 51.124 | 0.146
3$\times$2 | 2.279 | 2.202 | 52.155 | 0.139
RSD-C | 2$-$1 | 2.267 | 2.191 | 52.932 | 0.142
2$-$2 | 2.398 | 2.317 | 62.106 | 0.143
3$-$1 | 2.291 | 2.214 | 55.672 | 0.140
RSD-S | 2$\times$2 | 2.367 | 2.288 | 54.303 | 0.143
3$\times$2 | 2.432 | 2.350 | 62.251 | 0.140
3 | SD | 3 | 2.352 | 2.235 | 51.842 | 0.139
SpecTr | 3$\times$3 | 2.528 | 2.403 | 53.378 | 0.145
4$\times$3 | 2.554 | 2.428 | 54.173 | 0.139
RSD-C | 2$-$2$-$2 | 2.683 | 2.550 | 56.243 | 0.146
3$-$1$-$1 | 2.508 | 2.384 | 53.957 | 0.145
4$-$1$-$1 | 2.531 | 2.406 | 54.052 | 0.144
RSD-S | 3$\times$3 | 2.706 | 2.572 | 55.301 | 0.145
4$\times$3 | 2.734 | 2.599 | 62.092 | 0.140
4 | SD | 4 | 2.565 | 2.399 | 55.745 | 0.140
SpecTr | 5$\times$4 | 2.785 | 2.604 | 55.816 | 0.143
7$\times$4 | 2.761 | 2.582 | 53.047 | 0.137
RSD-C | 2$-$2$-$2$-$2 | 2.868 | 2.682 | 56.598 | 0.139
5$-$1$-$1$-$1 | 2.677 | 2.503 | 53.006 | 0.141
7$-$1$-$1$-$1 | 2.647 | 2.475 | 52.556 | 0.140
RSD-S | 5$\times$4 | 3.054 | 2.856 | 57.927 | 0.139
7$\times$4 | 3.059 | 2.860 | 56.124 | 0.144
5 | SD | 5 | 2.750 | 2.531 | 50.782 | 0.143
SpecTr | 6$\times$5 | 2.850 | 2.623 | 54.761 | 0.140
12$\times$5 | 2.730 | 2.512 | 48.408 | 0.138
RSD-C | 2$-$2$-$2$-$2$-$2 | 3.003 | 2.763 | 59.660 | 0.143
6$-$1$-$1$-$1$-$1 | 2.713 | 2.497 | 48.597 | 0.143
12$-$1$-$1$-$1$-$1 | 2.614 | 2.406 | 46.269 | 0.140
RSD-S | 6$\times$5 | 3.095 | 2.848 | 51.888 | 0.143
12$\times$5 | 3.081 | 2.835 | 51.306 | 0.140
Table 2: We summarize experiment results with Llama 2-7B target and 115M draft
for the WMT task with various draft lengths. Draft Length (DL) means the
(maximum) length for all draft token sequences generated by the draft model.
For decoders (Dec.), we consider Auto-Regressive sampling (AR), Speculative
Decoding (SD), SpecTr, Recursive Speculative Decoding with Constant branching
factors (RSD-C), Recursive Speculative Decoding with Stochastic Beam Search
(RSD-S). The contents in decoder specification (Spec.) have different meanings
for each decoder; $K\times L$ means the number $K$ of draft paths and draft
length $L$ for SpecTr; it describes constant branching factors for each level
of the tree (from root to leaf) for RSD-C; $K\times L$ means the beamwidth $K$
and draft length $L$ for RSD-S. Block efficiency (Eff.), Memory Bound Speed Up
(MBSU), Token Rate (TR), and Accuracy (Acc.) are given for each algorithm. For
each group of rows having the same complexity, we highlight the best values
for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-7B | WMT | 0 | AR | - | 1.000 | 1.000 | 37.631 | 0.374
2 | SD | 2 | 1.673 | 1.617 | 40.969 | 0.370
SpecTr | 2$\times$2 | 1.727 | 1.669 | 42.362 | 0.370
3$\times$2 | 1.757 | 1.698 | 43.190 | 0.376
RSD-C | 2$-$1 | 1.768 | 1.708 | 44.185 | 0.377
2$-$2 | 1.858 | 1.796 | 48.590 | 0.372
3$-$1 | 1.819 | 1.758 | 45.920 | 0.375
RSD-S | 2$\times$2 | 1.824 | 1.763 | 44.153 | 0.370
3$\times$2 | 1.912 | 1.847 | 45.014 | 0.373
3 | SD | 3 | 1.745 | 1.659 | 42.102 | 0.372
SpecTr | 3$\times$3 | 1.838 | 1.747 | 40.608 | 0.370
4$\times$3 | 1.863 | 1.771 | 41.160 | 0.372
RSD-C | 2$-$2$-$2 | 1.946 | 1.850 | 47.791 | 0.372
3$-$1$-$1 | 1.882 | 1.789 | 42.667 | 0.372
4$-$1$-$1 | 1.916 | 1.821 | 41.131 | 0.375
RSD-S | 3$\times$3 | 2.013 | 1.913 | 43.158 | 0.377
4$\times$3 | 2.070 | 1.968 | 47.280 | 0.373
4 | SD | 4 | 1.800 | 1.684 | 41.146 | 0.374
SpecTr | 5$\times$4 | 1.942 | 1.816 | 38.868 | 0.377
7$\times$4 | 1.991 | 1.862 | 40.518 | 0.378
RSD-C | 2$-$2$-$2$-$2 | 2.005 | 1.875 | 40.612 | 0.373
5$-$1$-$1$-$1 | 1.979 | 1.851 | 43.266 | 0.374
7$-$1$-$1$-$1 | 2.008 | 1.878 | 40.770 | 0.374
RSD-S | 5$\times$4 | 2.186 | 2.044 | 41.848 | 0.370
7$\times$4 | 2.258 | 2.111 | 45.639 | 0.375
5 | SD | 5 | 1.841 | 1.694 | 34.233 | 0.375
SpecTr | 6$\times$5 | 1.966 | 1.809 | 36.722 | 0.373
12$\times$5 | 1.977 | 1.819 | 36.132 | 0.375
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.039 | 1.877 | 38.175 | 0.375
6$-$1$-$1$-$1$-$1 | 2.002 | 1.842 | 40.785 | 0.372
12$-$1$-$1$-$1$-$1 | 2.047 | 1.883 | 37.631 | 0.377
RSD-S | 6$\times$5 | 2.249 | 2.069 | 39.514 | 0.379
12$\times$5 | 2.356 | 2.168 | 40.701 | 0.373
Table 3: We summarize experiment results with Llama 2-13B target and 115M
draft for the XSum task with various draft lengths. Draft Length (DL) means
the (maximum) length for all draft token sequences generated by the draft
model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-13B | XSum | 0 | AR | - | 1.000 | 1.000 | 28.141 | 0.166
2 | SD | 2 | 2.120 | 2.082 | 41.817 | 0.164
SpecTr | 2$\times$2 | 2.172 | 2.133 | 44.653 | 0.170
3$\times$2 | 2.212 | 2.173 | 42.566 | 0.165
RSD-C | 2$-$1 | 2.224 | 2.185 | 43.483 | 0.165
2$-$2 | 2.347 | 2.305 | 44.506 | 0.166
3$-$1 | 2.269 | 2.229 | 46.856 | 0.158
RSD-S | 2$\times$2 | 2.311 | 2.271 | 43.751 | 0.165
3$\times$2 | 2.412 | 2.370 | 45.704 | 0.162
3 | SD | 3 | 2.307 | 2.246 | 44.925 | 0.162
SpecTr | 3$\times$3 | 2.433 | 2.369 | 48.562 | 0.165
4$\times$3 | 2.455 | 2.391 | 43.066 | 0.167
RSD-C | 2$-$2$-$2 | 2.616 | 2.547 | 48.412 | 0.159
3$-$1$-$1 | 2.451 | 2.386 | 46.918 | 0.164
4$-$1$-$1 | 2.469 | 2.404 | 44.484 | 0.161
RSD-S | 3$\times$3 | 2.644 | 2.575 | 46.088 | 0.165
4$\times$3 | 2.713 | 2.642 | 48.354 | 0.170
4 | SD | 4 | 2.484 | 2.398 | 43.693 | 0.165
SpecTr | 5$\times$4 | 2.689 | 2.596 | 43.370 | 0.165
7$\times$4 | 2.671 | 2.578 | 41.911 | 0.164
RSD-C | 2$-$2$-$2$-$2 | 2.791 | 2.694 | 44.814 | 0.160
5$-$1$-$1$-$1 | 2.598 | 2.508 | 42.795 | 0.164
7$-$1$-$1$-$1 | 2.594 | 2.504 | 42.492 | 0.163
RSD-S | 5$\times$4 | 2.925 | 2.824 | 46.844 | 0.164
7$\times$4 | 2.981 | 2.878 | 45.746 | 0.164
5 | SD | 5 | 2.602 | 2.490 | 41.998 | 0.161
SpecTr | 6$\times$5 | 2.781 | 2.662 | 41.256 | 0.164
12$\times$5 | 2.678 | 2.563 | 39.336 | 0.159
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.922 | 2.797 | 44.311 | 0.168
6$-$1$-$1$-$1$-$1 | 2.648 | 2.535 | 40.041 | 0.163
12$-$1$-$1$-$1$-$1 | 2.553 | 2.444 | 38.607 | 0.161
RSD-S | 6$\times$5 | 3.030 | 2.899 | 47.832 | 0.161
12$\times$5 | 3.032 | 2.901 | 42.870 | 0.162
Table 4: We summarize experiment results with Llama 2-13B target and 115M
draft for the WMT task with various draft lengths. Draft Length (DL) means the
(maximum) length for all draft token sequences generated by the draft model.
For decoders (Dec.), we consider Auto-Regressive sampling (AR), Speculative
Decoding (SD), SpecTr, Recursive Speculative Decoding with Constant branching
factors (RSD-C), Recursive Speculative Decoding with Stochastic Beam Search
(RSD-S). The contents in decoder specification (Spec.) have different meanings
for each decoder; $K\times L$ means the number $K$ of draft paths and draft
length $L$ for SpecTr; it describes constant branching factors for each level
of the tree (from root to leaf) for RSD-C; $K\times L$ means the beamwidth $K$
and draft length $L$ for RSD-S. Block efficiency (Eff.), Memory Bound Speed Up
(MBSU), Token Rate (TR), and Accuracy (Acc.) are given for each algorithm. For
each group of rows having the same complexity, we highlight the best values
for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-13B | WMT | 0 | AR | - | 1.000 | 1.000 | 30.467 | 0.413
2 | SD | 2 | 1.662 | 1.632 | 38.250 | 0.410
SpecTr | 2$\times$2 | 1.717 | 1.686 | 36.007 | 0.405
3$\times$2 | 1.748 | 1.717 | 35.599 | 0.408
RSD-C | 2$-$1 | 1.760 | 1.729 | 36.755 | 0.405
2$-$2 | 1.852 | 1.819 | 37.389 | 0.407
3$-$1 | 1.815 | 1.783 | 37.600 | 0.408
RSD-S | 2$\times$2 | 1.810 | 1.778 | 36.507 | 0.404
3$\times$2 | 1.903 | 1.869 | 38.162 | 0.410
3 | SD | 3 | 1.738 | 1.692 | 34.550 | 0.410
SpecTr | 3$\times$3 | 1.828 | 1.780 | 33.706 | 0.409
4$\times$3 | 1.859 | 1.810 | 34.913 | 0.408
RSD-C | 2$-$2$-$2 | 1.946 | 1.895 | 38.604 | 0.411
3$-$1$-$1 | 1.884 | 1.834 | 35.514 | 0.409
4$-$1$-$1 | 1.910 | 1.860 | 35.173 | 0.408
RSD-S | 3$\times$3 | 1.997 | 1.945 | 36.920 | 0.409
4$\times$3 | 2.058 | 2.004 | 38.294 | 0.408
4 | SD | 4 | 1.793 | 1.731 | 31.712 | 0.408
SpecTr | 5$\times$4 | 1.927 | 1.860 | 33.392 | 0.408
7$\times$4 | 1.949 | 1.881 | 34.133 | 0.411
RSD-C | 2$-$2$-$2$-$2 | 1.997 | 1.927 | 34.046 | 0.409
5$-$1$-$1$-$1 | 1.974 | 1.905 | 34.561 | 0.405
7$-$1$-$1$-$1 | 1.999 | 1.930 | 34.209 | 0.405
RSD-S | 5$\times$4 | 2.183 | 2.107 | 35.807 | 0.412
7$\times$4 | 2.248 | 2.170 | 37.811 | 0.407
5 | SD | 5 | 1.827 | 1.748 | 33.288 | 0.406
SpecTr | 6$\times$5 | 1.949 | 1.865 | 31.442 | 0.409
12$\times$5 | 1.973 | 1.889 | 31.565 | 0.409
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.023 | 1.936 | 32.853 | 0.405
6$-$1$-$1$-$1$-$1 | 1.996 | 1.910 | 32.258 | 0.410
12$-$1$-$1$-$1$-$1 | 2.040 | 1.952 | 30.871 | 0.411
RSD-S | 6$\times$5 | 2.231 | 2.136 | 34.733 | 0.409
12$\times$5 | 2.342 | 2.241 | 35.124 | 0.407
Table 5: We summarize experiment results with Llama 2-70B target and 115M
draft for the XSum task with various draft lengths. Draft Length (DL) means
the (maximum) length for all draft token sequences generated by the draft
model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-70B | XSum | 0 | AR | - | 1.000 | 1.000 | 9.051 | 0.194
2 | SD | 2 | 2.103 | 2.096 | 15.005 | 0.188
SpecTr | 2$\times$2 | 2.164 | 2.157 | 15.188 | 0.189
3$\times$2 | 2.204 | 2.197 | 15.371 | 0.191
RSD-C | 2$-$1 | 2.189 | 2.181 | 15.294 | 0.187
2$-$2 | 2.322 | 2.314 | 16.121 | 0.189
3$-$1 | 2.239 | 2.231 | 15.497 | 0.197
RSD-S | 2$\times$2 | 2.288 | 2.280 | 15.798 | 0.189
3$\times$2 | 2.376 | 2.368 | 16.096 | 0.193
3 | SD | 3 | 2.263 | 2.251 | 15.391 | 0.189
SpecTr | 3$\times$3 | 2.410 | 2.397 | 15.920 | 0.195
4$\times$3 | 2.429 | 2.417 | 15.919 | 0.194
RSD-C | 2$-$2$-$2 | 2.566 | 2.553 | 16.741 | 0.197
3$-$1$-$1 | 2.393 | 2.381 | 15.926 | 0.198
4$-$1$-$1 | 2.421 | 2.409 | 15.969 | 0.189
RSD-S | 3$\times$3 | 2.603 | 2.590 | 16.982 | 0.190
4$\times$3 | 2.682 | 2.668 | 17.234 | 0.191
4 | SD | 4 | 2.442 | 2.425 | 15.860 | 0.191
SpecTr | 5$\times$4 | 2.613 | 2.595 | 16.190 | 0.193
7$\times$4 | 2.594 | 2.577 | 15.668 | 0.192
RSD-C | 2$-$2$-$2$-$2 | 2.730 | 2.712 | 16.335 | 0.191
5$-$1$-$1$-$1 | 2.561 | 2.544 | 15.826 | 0.193
7$-$1$-$1$-$1 | 2.585 | 2.568 | 15.610 | 0.186
RSD-S | 5$\times$4 | 2.880 | 2.861 | 17.159 | 0.194
7$\times$4 | 2.915 | 2.895 | 17.076 | 0.194
5 | SD | 5 | 2.520 | 2.499 | 16.058 | 0.188
SpecTr | 6$\times$5 | 2.658 | 2.636 | 15.437 | 0.194
12$\times$5 | 2.587 | 2.565 | 14.397 | 0.191
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.860 | 2.836 | 15.859 | 0.192
6$-$1$-$1$-$1$-$1 | 2.592 | 2.570 | 15.122 | 0.188
12$-$1$-$1$-$1$-$1 | 2.534 | 2.513 | 14.125 | 0.192
RSD-S | 6$\times$5 | 2.997 | 2.972 | 16.889 | 0.194
12$\times$5 | 2.981 | 2.956 | 16.040 | 0.195
Table 6: We summarize experiment results with Llama 2-70B target and 115M
draft for the WMT task with various draft lengths. Draft Length (DL) means the
(maximum) length for all draft token sequences generated by the draft model.
For decoders (Dec.), we consider Auto-Regressive sampling (AR), Speculative
Decoding (SD), SpecTr, Recursive Speculative Decoding with Constant branching
factors (RSD-C), Recursive Speculative Decoding with Stochastic Beam Search
(RSD-S). The contents in decoder specification (Spec.) have different meanings
for each decoder; $K\times L$ means the number $K$ of draft paths and draft
length $L$ for SpecTr; it describes constant branching factors for each level
of the tree (from root to leaf) for RSD-C; $K\times L$ means the beamwidth $K$
and draft length $L$ for RSD-S. Block efficiency (Eff.), Memory Bound Speed Up
(MBSU), Token Rate (TR), and Accuracy (Acc.) are given for each algorithm. For
each group of rows having the same complexity, we highlight the best values
for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-70B | WMT | 0 | AR | - | 1.000 | 1.000 | 9.710 | 0.439
2 | SD | 2 | 1.661 | 1.655 | 13.419 | 0.440
SpecTr | 2$\times$2 | 1.732 | 1.726 | 13.645 | 0.445
3$\times$2 | 1.756 | 1.750 | 13.911 | 0.445
RSD-C | 2$-$1 | 1.770 | 1.764 | 14.112 | 0.436
2$-$2 | 1.853 | 1.847 | 14.699 | 0.443
3$-$1 | 1.819 | 1.813 | 14.210 | 0.440
RSD-S | 2$\times$2 | 1.819 | 1.813 | 14.420 | 0.438
3$\times$2 | 1.907 | 1.900 | 14.758 | 0.439
3 | SD | 3 | 1.736 | 1.727 | 13.279 | 0.442
SpecTr | 3$\times$3 | 1.841 | 1.832 | 13.929 | 0.436
4$\times$3 | 1.861 | 1.852 | 13.778 | 0.437
RSD-C | 2$-$2$-$2 | 1.943 | 1.933 | 14.521 | 0.440
3$-$1$-$1 | 1.878 | 1.869 | 14.011 | 0.443
4$-$1$-$1 | 1.914 | 1.904 | 14.214 | 0.434
RSD-S | 3$\times$3 | 1.994 | 1.984 | 14.590 | 0.443
4$\times$3 | 2.069 | 2.058 | 14.937 | 0.441
4 | SD | 4 | 1.798 | 1.786 | 13.184 | 0.436
SpecTr | 5$\times$4 | 1.934 | 1.921 | 13.396 | 0.440
7$\times$4 | 1.955 | 1.942 | 13.263 | 0.436
RSD-C | 2$-$2$-$2$-$2 | 1.999 | 1.985 | 13.603 | 0.438
5$-$1$-$1$-$1 | 1.983 | 1.969 | 13.588 | 0.438
7$-$1$-$1$-$1 | 2.003 | 1.990 | 13.454 | 0.441
RSD-S | 5$\times$4 | 2.178 | 2.164 | 14.779 | 0.438
7$\times$4 | 2.255 | 2.240 | 15.028 | 0.439
5 | SD | 5 | 1.830 | 1.814 | 12.871 | 0.439
SpecTr | 6$\times$5 | 1.961 | 1.945 | 12.746 | 0.439
12$\times$5 | 1.970 | 1.954 | 12.432 | 0.442
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.029 | 2.012 | 12.815 | 0.443
6$-$1$-$1$-$1$-$1 | 2.008 | 1.991 | 13.097 | 0.440
12$-$1$-$1$-$1$-$1 | 2.046 | 2.029 | 12.729 | 0.438
RSD-S | 6$\times$5 | 2.240 | 2.221 | 14.213 | 0.439
12$\times$5 | 2.361 | 2.341 | 14.466 | 0.443
Table 7: We summarize experiment results with Llama 2-Chat-7B target and 115M
draft for the XSum task with various draft lengths. Draft Length (DL) means
the (maximum) length for all draft token sequences generated by the draft
model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-7B | XSum | 0 | AR | - | 1.000 | 1.000 | 41.651 | 0.092
2 | SD | 2 | 1.938 | 1.873 | 54.853 | 0.091
SpecTr | 2$\times$2 | 1.961 | 1.896 | 46.446 | 0.092
3$\times$2 | 1.972 | 1.906 | 46.728 | 0.092
RSD-C | 2$-$1 | 2.048 | 1.979 | 49.477 | 0.091
2$-$2 | 2.162 | 2.090 | 54.964 | 0.089
3$-$1 | 2.100 | 2.030 | 51.348 | 0.091
RSD-S | 2$\times$2 | 2.129 | 2.058 | 54.578 | 0.090
3$\times$2 | 2.220 | 2.146 | 52.765 | 0.090
3 | SD | 3 | 2.064 | 1.961 | 47.315 | 0.094
SpecTr | 3$\times$3 | 2.100 | 1.996 | 44.726 | 0.092
4$\times$3 | 2.103 | 1.999 | 45.956 | 0.088
RSD-C | 2$-$2$-$2 | 2.343 | 2.227 | 53.642 | 0.091
3$-$1$-$1 | 2.201 | 2.092 | 47.580 | 0.092
4$-$1$-$1 | 2.228 | 2.118 | 48.272 | 0.090
RSD-S | 3$\times$3 | 2.389 | 2.271 | 49.572 | 0.088
4$\times$3 | 2.469 | 2.347 | 55.614 | 0.092
4 | SD | 4 | 2.169 | 2.028 | 44.397 | 0.092
SpecTr | 5$\times$4 | 2.205 | 2.061 | 44.200 | 0.091
7$\times$4 | 2.190 | 2.047 | 42.661 | 0.092
RSD-C | 2$-$2$-$2$-$2 | 2.409 | 2.252 | 47.284 | 0.091
5$-$1$-$1$-$1 | 2.333 | 2.181 | 47.991 | 0.089
7$-$1$-$1$-$1 | 2.356 | 2.203 | 47.197 | 0.090
RSD-S | 5$\times$4 | 2.659 | 2.486 | 51.450 | 0.089
7$\times$4 | 2.710 | 2.534 | 51.212 | 0.093
5 | SD | 5 | 2.192 | 2.017 | 41.453 | 0.090
SpecTr | 6$\times$5 | 2.252 | 2.072 | 41.350 | 0.089
12$\times$5 | 2.171 | 1.998 | 38.959 | 0.090
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.477 | 2.279 | 50.571 | 0.090
6$-$1$-$1$-$1$-$1 | 2.368 | 2.179 | 42.483 | 0.091
12$-$1$-$1$-$1$-$1 | 2.370 | 2.181 | 42.743 | 0.090
RSD-S | 6$\times$5 | 2.720 | 2.503 | 47.129 | 0.090
12$\times$5 | 2.819 | 2.594 | 51.725 | 0.090
Table 8: We summarize experiment results with Llama 2-Chat-7B target and 115M
draft for the WMT task with various draft lengths. Draft Length (DL) means the
(maximum) length for all draft token sequences generated by the draft model.
For decoders (Dec.), we consider Auto-Regressive sampling (AR), Speculative
Decoding (SD), SpecTr, Recursive Speculative Decoding with Constant branching
factors (RSD-C), Recursive Speculative Decoding with Stochastic Beam Search
(RSD-S). The contents in decoder specification (Spec.) have different meanings
for each decoder; $K\times L$ means the number $K$ of draft paths and draft
length $L$ for SpecTr; it describes constant branching factors for each level
of the tree (from root to leaf) for RSD-C; $K\times L$ means the beamwidth $K$
and draft length $L$ for RSD-S. Block efficiency (Eff.), Memory Bound Speed Up
(MBSU), Token Rate (TR), and Accuracy (Acc.) are given for each algorithm. For
each group of rows having the same complexity, we highlight the best values
for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-7B | WMT | 0 | AR | - | 1.000 | 1.000 | 37.093 | 0.377
2 | SD | 2 | 1.639 | 1.584 | 41.655 | 0.379
SpecTr | 2$\times$2 | 1.664 | 1.608 | 40.765 | 0.379
3$\times$2 | 1.673 | 1.617 | 41.753 | 0.378
RSD-C | 2$-$1 | 1.739 | 1.681 | 42.819 | 0.378
2$-$2 | 1.813 | 1.752 | 45.130 | 0.375
3$-$1 | 1.784 | 1.724 | 43.227 | 0.378
RSD-S | 2$\times$2 | 1.786 | 1.726 | 43.396 | 0.378
3$\times$2 | 1.865 | 1.802 | 47.988 | 0.379
3 | SD | 3 | 1.701 | 1.617 | 39.019 | 0.378
SpecTr | 3$\times$3 | 1.739 | 1.653 | 42.205 | 0.380
4$\times$3 | 1.746 | 1.660 | 38.413 | 0.377
RSD-C | 2$-$2$-$2 | 1.903 | 1.809 | 41.284 | 0.375
3$-$1$-$1 | 1.854 | 1.762 | 40.906 | 0.379
4$-$1$-$1 | 1.877 | 1.784 | 40.827 | 0.379
RSD-S | 3$\times$3 | 1.945 | 1.849 | 44.995 | 0.376
4$\times$3 | 1.992 | 1.893 | 48.071 | 0.376
4 | SD | 4 | 1.758 | 1.644 | 39.401 | 0.375
SpecTr | 5$\times$4 | 1.802 | 1.685 | 35.691 | 0.378
7$\times$4 | 1.796 | 1.679 | 38.998 | 0.376
RSD-C | 2$-$2$-$2$-$2 | 1.954 | 1.827 | 39.405 | 0.375
5$-$1$-$1$-$1 | 1.930 | 1.804 | 40.046 | 0.376
7$-$1$-$1$-$1 | 1.951 | 1.824 | 42.786 | 0.379
RSD-S | 5$\times$4 | 2.094 | 1.958 | 40.487 | 0.377
7$\times$4 | 2.139 | 2.000 | 41.280 | 0.376
5 | SD | 5 | 1.791 | 1.648 | 33.954 | 0.381
SpecTr | 6$\times$5 | 1.824 | 1.678 | 37.011 | 0.379
12$\times$5 | 1.786 | 1.644 | 33.241 | 0.377
RSD-C | 2$-$2$-$2$-$2$-$2 | 1.992 | 1.833 | 40.157 | 0.377
6$-$1$-$1$-$1$-$1 | 1.957 | 1.801 | 39.944 | 0.376
12$-$1$-$1$-$1$-$1 | 1.987 | 1.829 | 36.625 | 0.376
RSD-S | 6$\times$5 | 2.153 | 1.981 | 37.368 | 0.377
12$\times$5 | 2.220 | 2.043 | 38.018 | 0.378
Table 9: We summarize experiment results with Llama 2-Chat-7B target and 115M
draft for the Dolly task with various draft lengths. Draft Length (DL) means
the (maximum) length for all draft token sequences generated by the draft
model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-7B | Dolly | 0 | AR | - | 1.000 | 1.000 | 37.596 | -
2 | SD | 2 | 2.122 | 2.051 | 50.489 | -
SpecTr | 2$\times$2 | 2.177 | 2.104 | 50.680 | -
3$\times$2 | 2.215 | 2.140 | 50.875 | -
RSD-C | 2$-$1 | 2.182 | 2.109 | 53.867 | -
2$-$2 | 2.253 | 2.178 | 52.362 | -
3$-$1 | 2.201 | 2.127 | 55.928 | -
RSD-S | 2$\times$2 | 2.239 | 2.164 | 56.598 | -
3$\times$2 | 2.278 | 2.202 | 52.466 | -
3 | SD | 3 | 2.168 | 2.061 | 50.839 | -
SpecTr | 3$\times$3 | 2.278 | 2.166 | 47.259 | -
4$\times$3 | 2.302 | 2.188 | 47.902 | -
RSD-C | 2$-$2$-$2 | 2.357 | 2.241 | 53.416 | -
3$-$1$-$1 | 2.252 | 2.141 | 47.475 | -
4$-$1$-$1 | 2.256 | 2.145 | 49.185 | -
RSD-S | 3$\times$3 | 2.382 | 2.264 | 48.472 | -
4$\times$3 | 2.409 | 2.290 | 50.245 | -
4 | SD | 4 | 2.258 | 2.112 | 44.857 | -
SpecTr | 5$\times$4 | 2.436 | 2.278 | 47.389 | -
7$\times$4 | 2.451 | 2.292 | 45.944 | -
RSD-C | 2$-$2$-$2$-$2 | 2.443 | 2.284 | 44.737 | -
5$-$1$-$1$-$1 | 2.368 | 2.215 | 46.960 | -
7$-$1$-$1$-$1 | 2.372 | 2.218 | 46.751 | -
RSD-S | 5$\times$4 | 2.521 | 2.357 | 44.591 | -
7$\times$4 | 2.547 | 2.382 | 45.908 | -
5 | SD | 5 | 2.257 | 2.076 | 40.617 | -
SpecTr | 6$\times$5 | 2.433 | 2.238 | 42.231 | -
12$\times$5 | 2.454 | 2.258 | 41.838 | -
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.461 | 2.265 | 41.159 | -
6$-$1$-$1$-$1$-$1 | 2.376 | 2.186 | 45.277 | -
12$-$1$-$1$-$1$-$1 | 2.371 | 2.182 | 41.597 | -
RSD-S | 6$\times$5 | 2.550 | 2.347 | 40.823 | -
12$\times$5 | 2.602 | 2.394 | 41.942 | -
Table 10: We summarize experiment results with Llama 2-Chat-13B target and
115M draft for the XSum task with various draft lengths. Draft Length (DL)
means the (maximum) length for all draft token sequences generated by the
draft model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-13B | XSum | 0 | AR | - | 1.000 | 1.000 | 28.727 | 0.112
2 | SD | 2 | 1.941 | 1.906 | 43.850 | 0.113
SpecTr | 2$\times$2 | 1.973 | 1.938 | 39.466 | 0.110
3$\times$2 | 1.976 | 1.941 | 38.639 | 0.114
RSD-C | 2$-$1 | 2.044 | 2.008 | 41.264 | 0.111
2$-$2 | 2.166 | 2.128 | 45.868 | 0.112
3$-$1 | 2.093 | 2.056 | 44.227 | 0.114
RSD-S | 2$\times$2 | 2.127 | 2.089 | 41.670 | 0.111
3$\times$2 | 2.232 | 2.193 | 46.500 | 0.111
3 | SD | 3 | 2.075 | 2.020 | 42.099 | 0.112
SpecTr | 3$\times$3 | 2.117 | 2.062 | 37.920 | 0.110
4$\times$3 | 2.127 | 2.071 | 37.372 | 0.110
RSD-C | 2$-$2$-$2 | 2.353 | 2.292 | 42.080 | 0.110
3$-$1$-$1 | 2.225 | 2.166 | 40.420 | 0.111
4$-$1$-$1 | 2.251 | 2.192 | 39.869 | 0.112
RSD-S | 3$\times$3 | 2.395 | 2.332 | 41.498 | 0.111
4$\times$3 | 2.481 | 2.416 | 42.730 | 0.110
4 | SD | 4 | 2.192 | 2.116 | 37.174 | 0.109
SpecTr | 5$\times$4 | 2.231 | 2.153 | 38.124 | 0.110
7$\times$4 | 2.231 | 2.153 | 36.645 | 0.111
RSD-C | 2$-$2$-$2$-$2 | 2.437 | 2.352 | 41.228 | 0.110
5$-$1$-$1$-$1 | 2.361 | 2.279 | 38.488 | 0.110
7$-$1$-$1$-$1 | 2.362 | 2.280 | 42.543 | 0.114
RSD-S | 5$\times$4 | 2.671 | 2.578 | 46.935 | 0.111
7$\times$4 | 2.752 | 2.657 | 42.186 | 0.108
5 | SD | 5 | 2.253 | 2.156 | 35.568 | 0.111
SpecTr | 6$\times$5 | 2.259 | 2.162 | 34.825 | 0.112
12$\times$5 | 2.201 | 2.107 | 33.259 | 0.111
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.519 | 2.411 | 38.598 | 0.112
6$-$1$-$1$-$1$-$1 | 2.385 | 2.283 | 36.961 | 0.110
12$-$1$-$1$-$1$-$1 | 2.385 | 2.282 | 36.284 | 0.109
RSD-S | 6$\times$5 | 2.749 | 2.631 | 43.266 | 0.111
12$\times$5 | 2.836 | 2.714 | 43.963 | 0.109
Table 11: We summarize experiment results with Llama 2-Chat-13B target and
115M draft for the WMT task with various draft lengths. Draft Length (DL)
means the (maximum) length for all draft token sequences generated by the
draft model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-13B | WMT | 0 | AR | - | 1.000 | 1.000 | 29.233 | 0.340
2 | SD | 2 | 1.729 | 1.699 | 39.502 | 0.346
SpecTr | 2$\times$2 | 1.806 | 1.774 | 38.003 | 0.338
3$\times$2 | 1.758 | 1.727 | 36.454 | 0.343
RSD-C | 2$-$1 | 1.758 | 1.727 | 36.551 | 0.357
2$-$2 | 1.928 | 1.894 | 39.740 | 0.345
3$-$1 | 1.891 | 1.858 | 38.799 | 0.344
RSD-S | 2$\times$2 | 1.956 | 1.922 | 39.189 | 0.333
3$\times$2 | 2.023 | 1.987 | 39.766 | 0.349
3 | SD | 3 | 1.774 | 1.727 | 34.434 | 0.349
SpecTr | 3$\times$3 | 1.899 | 1.849 | 35.247 | 0.342
4$\times$3 | 1.862 | 1.813 | 34.699 | 0.340
RSD-C | 2$-$2$-$2 | 2.037 | 1.984 | 41.584 | 0.341
3$-$1$-$1 | 1.928 | 1.877 | 35.565 | 0.351
4$-$1$-$1 | 1.880 | 1.831 | 34.997 | 0.353
RSD-S | 3$\times$3 | 1.997 | 1.945 | 36.662 | 0.348
4$\times$3 | 2.146 | 2.090 | 41.587 | 0.344
4 | SD | 4 | 1.895 | 1.829 | 33.551 | 0.347
SpecTr | 5$\times$4 | 1.964 | 1.896 | 34.044 | 0.352
7$\times$4 | 2.019 | 1.949 | 34.469 | 0.346
RSD-C | 2$-$2$-$2$-$2 | 2.155 | 2.080 | 37.331 | 0.352
5$-$1$-$1$-$1 | 2.183 | 2.107 | 38.244 | 0.334
7$-$1$-$1$-$1 | 2.080 | 2.008 | 40.071 | 0.350
RSD-S | 5$\times$4 | 2.297 | 2.218 | 38.790 | 0.345
7$\times$4 | 2.235 | 2.157 | 40.297 | 0.332
5 | SD | 5 | 1.969 | 1.885 | 33.378 | 0.352
SpecTr | 6$\times$5 | 2.070 | 1.981 | 33.048 | 0.339
12$\times$5 | 1.879 | 1.798 | 28.828 | 0.344
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.268 | 2.171 | 36.284 | 0.344
6$-$1$-$1$-$1$-$1 | 2.065 | 1.976 | 33.898 | 0.348
12$-$1$-$1$-$1$-$1 | 2.299 | 2.200 | 35.902 | 0.343
RSD-S | 6$\times$5 | 2.440 | 2.335 | 36.657 | 0.343
12$\times$5 | 2.369 | 2.267 | 34.976 | 0.352
Table 12: We summarize experiment results with Llama 2-Chat-13B target and
115M draft for the Dolly task with various draft lengths. Draft Length (DL)
means the (maximum) length for all draft token sequences generated by the
draft model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-13B | Dolly | 0 | AR | - | 1.000 | 1.000 | 29.672 | -
2 | SD | 2 | 2.103 | 2.066 | 41.930 | -
SpecTr | 2$\times$2 | 2.158 | 2.120 | 45.409 | -
3$\times$2 | 2.187 | 2.148 | 45.432 | -
RSD-C | 2$-$1 | 2.163 | 2.125 | 44.332 | -
2$-$2 | 2.241 | 2.201 | 43.922 | -
3$-$1 | 2.181 | 2.143 | 43.922 | -
RSD-S | 2$\times$2 | 2.230 | 2.191 | 48.240 | -
3$\times$2 | 2.262 | 2.222 | 43.045 | -
3 | SD | 3 | 2.139 | 2.083 | 40.211 | -
SpecTr | 3$\times$3 | 2.257 | 2.197 | 39.783 | -
4$\times$3 | 2.280 | 2.220 | 45.590 | -
RSD-C | 2$-$2$-$2 | 2.329 | 2.268 | 41.478 | -
3$-$1$-$1 | 2.225 | 2.167 | 40.625 | -
4$-$1$-$1 | 2.241 | 2.182 | 42.782 | -
RSD-S | 3$\times$3 | 2.369 | 2.307 | 40.248 | -
4$\times$3 | 2.387 | 2.324 | 41.529 | -
4 | SD | 4 | 2.235 | 2.158 | 38.628 | -
SpecTr | 5$\times$4 | 2.407 | 2.324 | 41.825 | -
7$\times$4 | 2.431 | 2.347 | 41.870 | -
RSD-C | 2$-$2$-$2$-$2 | 2.415 | 2.331 | 40.088 | -
5$-$1$-$1$-$1 | 2.338 | 2.257 | 38.123 | -
7$-$1$-$1$-$1 | 2.344 | 2.263 | 39.884 | -
RSD-S | 5$\times$4 | 2.501 | 2.415 | 38.642 | -
7$\times$4 | 2.528 | 2.440 | 40.117 | -
5 | SD | 5 | 2.223 | 2.127 | 35.653 | -
SpecTr | 6$\times$5 | 2.401 | 2.298 | 37.463 | -
12$\times$5 | 2.426 | 2.322 | 35.909 | -
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.429 | 2.325 | 37.942 | -
6$-$1$-$1$-$1$-$1 | 2.342 | 2.242 | 35.915 | -
12$-$1$-$1$-$1$-$1 | 2.348 | 2.247 | 35.768 | -
RSD-S | 6$\times$5 | 2.532 | 2.424 | 35.828 | -
12$\times$5 | 2.578 | 2.468 | 36.071 | -
Table 13: We summarize experiment results with Llama 2-Chat-70B target and
115M draft for the XSum task with various draft lengths. Draft Length (DL)
means the (maximum) length for all draft token sequences generated by the
draft model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-70B | XSum | 0 | AR | - | 1.000 | 1.000 | 9.183 | 0.118
2 | SD | 2 | 1.905 | 1.899 | 14.110 | 0.121
SpecTr | 2$\times$2 | 1.933 | 1.926 | 14.090 | 0.121
3$\times$2 | 1.939 | 1.932 | 14.079 | 0.122
RSD-C | 2$-$1 | 2.017 | 2.010 | 14.683 | 0.118
2$-$2 | 2.130 | 2.123 | 15.293 | 0.118
3$-$1 | 2.074 | 2.067 | 15.009 | 0.118
RSD-S | 2$\times$2 | 2.093 | 2.086 | 14.958 | 0.119
3$\times$2 | 2.195 | 2.188 | 15.856 | 0.119
3 | SD | 3 | 2.018 | 2.008 | 14.162 | 0.122
SpecTr | 3$\times$3 | 2.077 | 2.067 | 14.294 | 0.118
4$\times$3 | 2.072 | 2.061 | 14.279 | 0.123
RSD-C | 2$-$2$-$2 | 2.284 | 2.272 | 15.483 | 0.120
3$-$1$-$1 | 2.175 | 2.164 | 14.922 | 0.119
4$-$1$-$1 | 2.210 | 2.199 | 15.043 | 0.119
RSD-S | 3$\times$3 | 2.357 | 2.345 | 15.857 | 0.118
4$\times$3 | 2.443 | 2.431 | 16.285 | 0.120
4 | SD | 4 | 2.114 | 2.100 | 14.447 | 0.122
SpecTr | 5$\times$4 | 2.169 | 2.154 | 13.991 | 0.121
7$\times$4 | 2.159 | 2.144 | 13.786 | 0.125
RSD-C | 2$-$2$-$2$-$2 | 2.364 | 2.349 | 15.073 | 0.121
5$-$1$-$1$-$1 | 2.307 | 2.291 | 14.674 | 0.116
7$-$1$-$1$-$1 | 2.324 | 2.308 | 14.713 | 0.120
RSD-S | 5$\times$4 | 2.610 | 2.593 | 16.231 | 0.121
7$\times$4 | 2.684 | 2.666 | 16.446 | 0.118
5 | SD | 5 | 2.169 | 2.151 | 14.033 | 0.121
SpecTr | 6$\times$5 | 2.200 | 2.181 | 13.537 | 0.120
12$\times$5 | 2.152 | 2.134 | 12.561 | 0.122
RSD-C | 2$-$2$-$2$-$2$-$2 | 2.447 | 2.427 | 14.116 | 0.120
6$-$1$-$1$-$1$-$1 | 2.331 | 2.311 | 13.994 | 0.117
12$-$1$-$1$-$1$-$1 | 2.330 | 2.311 | 13.503 | 0.121
RSD-S | 6$\times$5 | 2.690 | 2.667 | 15.870 | 0.120
12$\times$5 | 2.786 | 2.762 | 15.486 | 0.119
Table 14: We summarize experiment results with Llama 2-Chat-70B target and
115M draft for the WMT task with various draft lengths. Draft Length (DL)
means the (maximum) length for all draft token sequences generated by the
draft model. For decoders (Dec.), we consider Auto-Regressive sampling (AR),
Speculative Decoding (SD), SpecTr, Recursive Speculative Decoding with
Constant branching factors (RSD-C), Recursive Speculative Decoding with
Stochastic Beam Search (RSD-S). The contents in decoder specification (Spec.)
have different meanings for each decoder; $K\times L$ means the number $K$ of
draft paths and draft length $L$ for SpecTr; it describes constant branching
factors for each level of the tree (from root to leaf) for RSD-C; $K\times L$
means the beamwidth $K$ and draft length $L$ for RSD-S. Block efficiency
(Eff.), Memory Bound Speed Up (MBSU), Token Rate (TR), and Accuracy (Acc.) are
given for each algorithm. For each group of rows having the same complexity,
we highlight the best values for all columns except accuracy.
|  |  |  |  | Eff. | MBSU | TR | Acc.
---|---|---|---|---|---|---|---|---
Model | Task | DL | Dec. | Spec. |  |  |  |
Llama 2-Chat-70B | WMT | 0 | AR | - | 1.000 | 1.000 | 9.731 | 0.426
2 | SD | 2 | 1.647 | 1.642 | 13.286 | 0.424
SpecTr | 2$\times$2 | 1.668 | 1.663 | 13.389 | 0.425
3$\times$2 | 1.680 | 1.674 | 13.212 | 0.423
RSD-C | 2$-$1 | 1.738 | 1.732 | 13.740 | 0.422
2$-$2 | 1.819 | 1.813 | 14.279 | 0.423
3$-$1 | 1.790 | 1.783 | 14.074 | 0.422
RSD-S | 2$\times$2 | 1.790 | 1.784 | 14.025 | 0.425
3$\times$2 | 1.871 | 1.865 | 14.557 | 0.423
3 | SD | 3 | 1.709 | 1.701 | 13.096 | 0.422
SpecTr | 3$\times$3 | 1.742 | 1.733 | 13.094 | 0.427
4$\times$3 | 1.749 | 1.740 | 13.029 | 0.425
RSD-C | 2$-$2$-$2 | 1.913 | 1.903 | 14.173 | 0.426
3$-$1$-$1 | 1.857 | 1.848 | 13.773 | 0.424
4$-$1$-$1 | 1.888 | 1.878 | 14.000 | 0.425
RSD-S | 3$\times$3 | 1.946 | 1.937 | 14.392 | 0.425
4$\times$3 | 2.007 | 1.997 | 14.908 | 0.424
4 | SD | 4 | 1.760 | 1.748 | 13.020 | 0.425
SpecTr | 5$\times$4 | 1.804 | 1.792 | 12.567 | 0.424
7$\times$4 | 1.796 | 1.784 | 12.320 | 0.426