We present speculative sampling, an algorithm for accelerating transformer
decoding by enabling the generation of multiple tokens from each transformer
call. Our algorithm relies on the observation that the latency of parallel
scoring of short continuations, generated by a faster but less powerful draft
model, is comparable to that of sampling a single token from the larger target
model. This is combined with a novel modified rejection sampling scheme which
preserves the distribution of the target model within hardware numerics. We
benchmark speculative sampling with Chinchilla, a 70 billion parameter language
model, achieving a 2-2.5x decoding speedup in a distributed setup, without
compromising the sample quality or making modifications to the model itself.
Charlie Chen Sebastian Borgeaud Geoffrey Irving Jean-Baptiste Lespiau Laurent Sifre John Jumper
2023-2-3
Accelerating Large Language Model Decoding
with Speculative Sampling
Charlie Chen1, Sebastian Borgeaud1, Geoﬀrey Irving1, Jean-Baptiste Lespiau1, Laurent Sifre1and John
Jumper1
1All authors from DeepMind
We present speculative sampling, an algorithm for accelerating transformer decoding by enabling the
generation of multiple tokens from each transformer call. Our algorithm relies on the observation that
the latency of parallel scoring of short continuations, generated by a faster but less powerful draft
model, is comparable to that of sampling a single token from the larger target model. This is combined
with a novel modiﬁed rejection sampling scheme which preserves the distribution of the target model
withinhardwarenumerics. WebenchmarkspeculativesamplingwithChinchilla,a70billionparameter
language model, achieving a 2–25decoding speedup in a distributed setup, without compromising
the sample quality or making modiﬁcations to the model itself.
Introduction
Scaling transformer models to 500B+ parameters has led to large performance improvements on
many natural language, computer vision and reinforcement learning tasks (Arnab et al., 2021; Brown
etal.,2020;Chowdheryetal.,2022;Dosovitskiyetal.,2020;Hoﬀmannetal.,2022;Raeetal.,2021).
However, transformer decoding remains a highly costly and ineﬃcient process in this regime.
Transformer sampling is typically memory bandwidth bound (Shazeer, 2019), so for a given set of
hardware, the time to generate a single token in transformer models is proportional to a ﬁrst order
approximationtothesizeofparametersandthesizeofthetransformermemory. Thesizeoflanguage
models also necessitates serving with model parallelism – adding communication overheads (Pope
et al., 2022) and multiplying resource requirements. Since each new token depends on the past,
many such transformer calls are required to sample a new sequence.
We present an algorithm to accelerate transformer sampling for latency critical applications, which
we call speculative sampling (SpS). This is achieved by:
1.Generating a short draft of length 𝐾. This can be attained with either a parallel model (Stern
et al., 2018) or by calling a faster, auto-regressive model 𝐾times. We shall refer to this model
as thedraft model , and focus on the case where it is auto-regressive.
2.Scoring the draft using the larger, more powerful model from we wish to sample from. We shall
refer to this model as the target model .
3.Using a modiﬁed rejection sampling scheme, accept a subset of the 𝐾draft tokens from left to
right, recovering the distribution of the target model in the process.
Intuitively, there are often sequences where the next token might be “obvious”. Therefore, if there is
strongagreementbetweenthedraftandtargetmodel’sdistributionsonagiventokenorsub-sequence
of tokens, this setup permits the generation of multiple tokens each time the target model is called .
We show that the expected acceptance rate of draft tokens is suﬃcient to oﬀset the overhead of the
Corresponding author(s): ccharlie@deepmind.com
©2023 DeepMind. All rights reservedarXiv:2302.01318v1  [cs.CL]  2 Feb 2023
Accelerating Large Language Model Decoding with Speculative Sampling
draftingprocessforlargelanguagemodels,resultinginaneﬀectiveandpracticalmethodforreducing
sampling latency without the need for modifying the target model or biasing the sample distribution.
Depending on the evaluation domain, SpS leads to a 2–25speedup when sampling from Chinchilla
(Hoﬀmann et al., 2022). Notably, the mean tokens per second with SpS often exceeds the idealised
ceiling on auto-regressive sampling speed imposed by the memory bandwidth.
Related Work
There has been a substantial body of work focused on improving sampling latency of large transform-
ers and other auto-regressive models.
Since sampling performance is heavily coupled with the model size in memory, quantisation to
int8or even int4(Dettmers et al., 2022; Yao et al., 2022) and distillation (Jiao et al., 2020; Sanh
et al., 2019) of transformers are eﬀective techniques for reducing sampling latency with little to no
performance penalty. The observation that model size contributes less to the ﬁnal performance than
expected (Hoﬀmann et al., 2022) has also encouraged smaller language models in general.
During sampling, a cache of the keys and values is maintained for every attention layer, and could
become a memory bandwidth bottleneck as the batch size increases. Methods such as multi-query
attention (Shazeer, 2019) aims to improve sampling performance by shrinking this cache. However
thesetechniquesaremosteﬀectiveatmaximisingthroughout(atlargerbatchsizes)insteadoflatency,
especially for larger models where the majority of the memory bandwidth budget is consumed by the
parameters.
Using a combination of the above techniques, in addition to a number of low-level optimisations to
TPUs, Pope et al. (2022) have greatly improved the serving latency and eﬃciency of PaLM 540B.
There is an existing body of similar work exploiting the eﬃciency of transformers and sequence
models operating in parallel. This includes block parallel sampling (Stern et al., 2018), aggressive
decoding (Ge et al., 2022), in addition to some work in parallelizing autoregressive models in the
image domain (Song et al., 2021; Wiggers and Hoogeboom, 2020). These methods have yet to be
adapted to typical language model use-cases since they either only work with greedy sampling, bias
the results or are focused on other modalities. Further, to our knowledge none of these techniques
have been scaled to distributed setups, which is necessary for the most expensive decoders with the
tens or hundreds of billions of parameters.
Coincidentally, the work in this manuscript was undertaken concurrently and independently of
the work on speculative decoding from Leviathan et al. (2022). We focus more heavily the distributed
serving setting for large models and oﬀer some incremental optimisations, but otherwise the core
underlying idea is the same.
Auto-regressive Sampling
Whilst transformers can be trained eﬃciently and in parallel on TPUs and GPUs, samples are typically
drawn auto-regressively (See algorithm 1). For most applications, auto-regressive sampling (ArS) is
highlymemorybandwidthboundandthuscannotmakeeﬀectiveuseofmodernacceleratorhardware
(Shazeer, 2019). A memory bound model call only generates a single token for every sequence in the
batch, hence generating multiple tokens introduces a large amount of latency in any system which
makes use of it.
2
Accelerating Large Language Model Decoding with Speculative Sampling
Thisisespeciallyproblematicasthenumberofparameters inthemodelincreases. Sinceallthemodel
parameters need to pass through at least one accelerator chip, the model size divided by the total
memory bandwidth across all chips gives us a hard ceiling on the maximum auto-regressive sampling
speed. Larger models also require serving on multiple accelerators, introducing a further source of
latency due to inter-device communication overheads.
Algorithm 1 Auto-regressive (ArS) with Auto-Regressive Models
Givenauto-regressivetargetmodel 𝑞¹jºandinitialpromptsequence 𝑥1𝑥𝑡andtargetsequence
length𝑇.
Initialise𝑛 𝑡.
while𝑛 𝑇do
Sample𝑥𝑛¸1𝑞¹𝑥j𝑥1𝑥𝑛º
𝑛 𝑛¸1
end while
Algorithm 2 Speculative Sampling (SpS) with Auto-Regressive Target and Draft Models
Given lookahead 𝐾and minimum target sequence length 𝑇.
Given auto-regressive target model 𝑞¹jº, and auto-regressive draft model 𝑝¹jº, initial prompt
sequence𝑥0𝑥𝑡.
Initialise𝑛 𝑡.
while𝑛 𝑇do
for𝑡=1 :𝐾do
Sample draft auto-regressively ˜𝑥𝑡𝑝¹𝑥j𝑥1𝑥𝑛˜𝑥1 ˜𝑥𝑡1º
end for
In parallel, compute 𝐾¸1sets of logits from drafts ˜𝑥1 ˜𝑥𝐾:
𝑞¹𝑥j𝑥1𝑥𝑛º 𝑞¹𝑥j𝑥1𝑥𝑛˜𝑥1º  𝑞¹𝑥j𝑥1𝑥𝑛˜𝑥1 ˜𝑥𝐾º
for𝑡=1 :𝐾do
Sample𝑟𝑈»01¼from a uniform distribution.
if𝑟 min
1𝑞¹𝑥j𝑥1𝑥𝑛¸𝑡1º
𝑝¹𝑥j𝑥1𝑥𝑛¸𝑡1º
,then
Set𝑥𝑛¸𝑡 ˜𝑥𝑡and𝑛 𝑛¸1.
else
sample𝑥𝑛¸𝑡¹𝑞¹𝑥j𝑥1𝑥𝑛¸𝑡1º𝑝¹𝑥j𝑥1𝑥𝑛¸𝑡1ºº¸and exit for loop.
end if
end for
If all tokens 𝑥𝑛¸1𝑥𝑛¸𝐾are accepted, sample extra token 𝑥𝑛¸𝐾¸1𝑞¹𝑥j𝑥1𝑥𝑛𝑥𝑛¸𝐾ºand
set𝑛 𝑛¸1.
end while
Speculative Sampling
Conditional Scoring
For speculative sampling (See algorithm 2), we ﬁrst make the observation that computing the logits
of a short continuation of 𝐾tokens in parallel has a very similar latency to that of sampling a single
3
Accelerating Large Language Model Decoding with Speculative Sampling
token. We focus our attention on large transformers, sharded in the Megatron style (Shoeybi et al.,
2019). For these models the majority of sampling time can be attributed to three components:
1.Linear Layers: For small batch sizes, each linear layer only processes a small number of
embeddings. This causes the dense matrix multiplies in the feed-forward layers, queries, keys,
values computations and the ﬁnal attention projection to become memory bound. For small 𝐾,
this will continue to be memory bound and therefore take a similar amount of time.
2.The Attention Mechanism: Theattentionmechanismisalsomemorybound. Duringsampling,
we maintain a cache of all the keys and values of the previous tokens in the sequence to avoid
re-computation. These KV-caches are large, and accounts for the majority of the memory
bandwidth utilisation for the attention mechanism. However, since the KV-cache size does not
change as we increase 𝐾, there is little to no delta in this component.
3.All-reduces: Asmodelsgrowinsize, itsparametersneedtobedividedacrossmultipleaccelera-
tors, leading to communication overheads. With Megatron, this manifests itself as an all-reduce
after every feed-forward and attention layer. Since only the activations for a small number of
tokens are transmitted, this operation is typically latency bound instead of throughput bound
for both sampling and scoring (for small 𝐾). Again, this results in a similar amount of time
spent in the two cases.
Other sources of overhead may exist, depending on the exact transformer implementation. Therefore
it is still possible that the choice of positioning encoding, decoding method (e.g. a sort might be
required for nucleus sampling), hardware limitations etc. can introduce some deltas between scoring
and sampling. However, if the conditions are met such that the above components dominate then
scoring should not be signiﬁcantly slower for small 𝐾.
Modiﬁed Rejection Sampling
We require a method to recover the distribution of the target model from samples from the draft
model, and logits of said tokens from both models.
To achieve this, we introduce the following rejection sampling scheme of the drafted tokens. Given a
sequence of tokens 𝑥1𝑥𝑛, and𝐾draft tokens ˜𝑥𝑛¸1 ˜𝑥𝑛¸𝐾generated from 𝑝¹jº, we accept ˜𝑥𝑛¸1
with probability:
min
1𝑞¹˜𝑥𝑛¸1j𝑥1𝑥𝑛º
𝑝¹˜𝑥𝑛¸1j𝑥1𝑥𝑛º
Where𝑞¹˜𝑥𝑛¸1j𝑥1𝑥𝑛ºand𝑝¹˜𝑥𝑛¸1j𝑥1𝑥𝑛ºare the probability of ˜𝑥𝑛¸1according to the target and
draft models respectively, conditioned on the context so far.
If the token is accepted, we set 𝑥𝑛¸1 ˜𝑥𝑛¸1and repeat the process for ˜𝑥𝑛¸2until either a token
is rejected or all tokens have been accepted.
If˜𝑥𝑛¸1is rejected, we resample 𝑥𝑛¸1from the following distribution:
𝑥𝑛¸1¹𝑞¹𝑥j𝑥1𝑥𝑛º𝑝¹𝑥j𝑥1𝑥𝑛ºº¸
Where¹º¸denotes:
¹𝑓¹𝑥ºº¸=max¹0 𝑓¹𝑥ººÍ
𝑥max¹0 𝑓¹𝑥ºº
By applying this sequentially, we recover the distribution of the target model for the accepted tokens
(see proof in Theorem 1) within hardware numerics. Note that:
4
Accelerating Large Language Model Decoding with Speculative Sampling
•At least one token will always be generated from a draft-accept loop – if the ﬁrst token is
rejected, a valid token is resampled.
•Since the ﬁnal token of the draft gives us the logits for the next token, if every drafted token is
accepted, we can sample from it normally. This gives us a maximum of 𝐾¸1tokens per loop,
over the naive implementation which would only return 𝐾tokens.
With standard sampling methods such as nucleus, top-k sampling and adjusting temperature, we
can modify the probabilities accordingly before applying this rejection sampling scheme. We have
observed that the overall acceptance rate is robust to the exact parameters used.
Because we do not interact with the body of the transformer itself, this method can be used in
conjunction many other techniques for accelerating or optimising the memory use of sampling, such
as quantisation and multi-query attention.
Choice of Draft Models
Since the acceptance criterion guarantees the distribution of the target model in our samples, we are
free to choose the method for drafting a continuation as long as it exposes logits, and there is a high
enough acceptance rate and/or low enough latency to break-even. There exist several approaches
here:
•Incorporating draft generation into the target model, and train the model from the start. This
is the strategy used by Stern et al. (2018), which adds multiple heads into the transformer to
generate multiple tokens.
•Using sequence level distillation (Kim and Rush, 2016) to generate a second model which
predicts𝐾tokens in parallel. This strategy was employed by Ge et al. (2022).
•Set a portion of the activations of the target model as an input to the draft model, and train the
draft model with this input.
Although these methods will likely yield powerful drafts, they require a large number of data gener-
ated from the target model or changes to the target model. Sequence level distillation in particular
would require a large compute budget. This makes them less practical for large scale applications.
Whilst large language models produce better samples, intuitively there are "easier" tokens to predict
for which smaller models may be suﬃcient. Therefore we may simply use a smaller version of the
target language model as the draft and obtain high acceptance rates. This would also be convenient
from an engineering and workﬂow perspective, since robust tooling for such models should already
exist to train the target model in the ﬁrst place.
Results
We train a 4 billion parameter draft model optimised for sampling latency on 16 TPU v4s – the same
hardware that is typically used to serve Chinchilla for research purposes. This model was trained
with the same tokeniser and dataset as Chinchilla, with a slightly smaller width and with only 8
layers. The relatively few number of layers allows it to achieve a sampling speed of 1.8ms/token
comparedto 14.1ms/token forChinchilla. Fordetails,pleaserefertothehyperparametersinTable2.
For distributed setups it is insuﬃcient to naively choose a small model as the draft, since diﬀer-
ent models have diﬀerent optimal inference setups. For example, it is typical to serve Chinchilla 70B
5
Accelerating Large Language Model Decoding with Speculative Sampling
Table1jChinchillaperformanceandspeedonXSumandHumanEvalwithnaiveandspeculative
sampling at batch size 1 and 𝐾=4. XSum was executed with nucleus parameter 𝑝=08, and
HumanEval with 𝑝=095and temperature 08.
Sampling Method Benchmark Result Mean Token Time Speed Up
ArS (Nucleus)XSum (ROUGE-2)0.112 14.1ms/Token 1
SpS (Nucleus) 0.114 7.52ms/Token 192
ArS (Greedy)XSum (ROUGE-2)0.157 14.1ms/Token 1
SpS (Greedy) 0.156 7.00ms/Token 201
ArS (Nucleus)HumanEval (100 Shot)45.1% 14.1ms/Token 1
SpS (Nucleus) 47.0% 5.73ms/Token 246
on 16 TPU v4s (where it achieves the aforementioned 14.1ms/token ), whereas a chinchilla-optimal
7B achieves its lowest sampling latency on 4 TPU v4s (where it achieves 5ms/token ). For smaller
models, the additional memory bandwidth and ﬂops are insuﬃcient to oﬀset the additional communi-
cation overhead between more devices – serving a 7B on 16 TPUs actually increases the latency. This
means the 7B would provide only a modest speedup if used as a draft with its optimal topology, and
we will not make full utilisation of the hardware during drafting.
We can sidestep this issue by training a wider model with a relatively few number of layers in
order to minimise communication overhead. It has been observed that the performance of language
models is relatively robust to changes in model aspect ratio (Levine et al., 2020), so this allows us
to serve a powerful draft model which can be sampled rapidly on the same hardware as the target
model.
Evaluation on XSum and HumanEval
We evaluate speculative sampling with Chinchilla on two tasks and summarize the results in Table 1:
•The XSum (Narayan et al., 2018) benchmark. This is a natural language summarisation task
using a 1-shot prompt where we sample a total of 11,305 sequences with a maximum sequence
length 128.
•The 100-shot HumanEval task (Chen et al., 2021). This is a code generation task involves the
generation of 16,400 samples with a maximum sequence length of 512.
Even with greedy sampling, a single token deviating due to numerics could result in two sequences
diverging wildly. Since pseudo-random seeds are processed diﬀerently between ArS and SpS, and
because the diﬀerent computation graphs lead to diﬀerent numerics, we cannot not expect identical
outputs. However, we expect the samples to come from the same distribution within numerics and
we empirically verify this by evaluating these benchmarks.
We run the tasks at batch size 1 with SpS and ArS. The time taken per SpS/ArS loop has low
variance, and we can measure it directly from TPU proﬁles. To obtain the average speedup, standard
deviations and other metrics, we log the amount of tokens generated for each speculative loop. In
Table 1 we show the performance on the XSum and HumanEval benchmarks for naive and speculative
sampling with Chinchilla.
6
Accelerating Large Language Model Decoding with Speculative Sampling
We obtain a substantial speedup in both tasks, with HumanEval reaching speedups of almost 25.
Yet, we have parity in the benchmark metrics – the underlying samples distribution is provably the
sameuptonumerics, andthisveriﬁesthatthedraftmodelisnotbiasingtheresultsempirically. Inthe
case of HumanEval and greedy XSum, this speedup exceeded the theoretical memory bandwidth limit
of the hardware for autoregressive sampling (model size divided by the total memory bandwidth).
Acceptance rate changes per domain
It is apparent that the acceptance rate is dependent on the application and the decoding method.
HumanEvalachievesasigniﬁcantlylargerspeedup—Wehypothesizethatthisisduetoacombination
of code containing a lot of common sub-sequences (e.g. for i in range(len(arr)): would be
relatively easy for a draft model to guess), is often decomposed into a smaller set of shorter tokens
and the temperature value sharpening both the draft and target logits.
01234567
Number of draft tokens (K)40060080010001200140016001800ms
Mean Sampling Time (128 tokens)
Human Eval
XSum
01234567
Number of draft tokens (K)0.50.60.70.80.91.0
Acceptance rate
Human Eval
XSum
01234567
Number of draft tokens (K)1416182022242628ms
Total loop time
Figure 1jLeft:The average time to generate 128 tokens, with standard deviation. Note that as 𝐾
increases, the overall speedup plateaus or even regresses, with XSum being optimal at 𝐾=3. The
variance consistently increases with 𝐾.Middle:The average number of tokens accepted divided by
𝐾¸1– this serves as a measure of the overall eﬃciency of the modiﬁed rejection scheme, which
decreases with the lookahead. Right:Average time per loop increases approximately linearly with 𝐾
duetotheincreasednumberofmodelcalls. Notethatthegradientisslightlyhigherthanthesampling
speed of the draft model, due to additional overheads in nucleus decoding.
Trade oﬀ between longer drafts and more frequent scoring
Wevisualisethetrade-oﬀofincreasing 𝐾,thenumberoftokenssampledbythedraftmodelinFigure1.
As𝐾increases, we need fewer scoring calls from the large models to generate the same sequence
length, potentially giving us a larger speedup. However, the total loop time increases approximately
linearly with the larger number of draft model calls and small increases in the scoring time. The
overall eﬃciency of the proportion of accepted tokens decreases as 𝐾increases, since later tokens
depend on the acceptance of previous tokens. This results in the average speedup plateauing or
even degrading with a larger 𝐾(for example, XSum with nucleus’s latency is minimised at 𝐾=3),
depending on the domain.
7
Accelerating Large Language Model Decoding with Speculative Sampling
Further, even though larger values of 𝐾may yield marginally greater mean speedups in certain
circumstances, it also increases variance of the time to generate a full sequence. This could be
problematic for settings where the P90, P99 latencies of concern.
Conclusion
Inthiswork,wedemonstrateanewalgorithmandworkﬂowforacceleratingthedecodingoflanguage
models. Speculative sampling does not require making any modiﬁcations to the target language
model’s parameters or architecture, is provably lossless within numerics, scales well with the appro-
priate draft model and complements many existing techniques for reducing latency in the small batch
size setting.
We optimise and scale the technique to Chinchilla 70B using a draft model which was easy to
trainwithexistinginfrastructure, demonstratingthatityieldsalargespeedupacrossbenchmarktasks
and common decoding methods in the process. We verify that it is indeed lossless empirically in its
downstream tasks.
References
A.Arnab,M.Dehghani,G.Heigold,C.Sun,M.Lucic,andC.Schmid. Vivit: Avideovisiontransformer.
In2021 IEEE/CVF International Conference on Computer Vision (ICCV) , pages 6816–6826. IEEE
Computer Society, 2021.
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam,
G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information
processing systems , 33:1877–1901, 2020.
M.Chen,J.Tworek,H.Jun,Q.Yuan,H.P.deOliveiraPinto,J.Kaplan,H.Edwards,Y.Burda,N.Joseph,
G.Brockman,A.Ray,R.Puri,G.Krueger,M.Petrov,H.Khlaaf,G.Sastry,P.Mishkin,B.Chan,S.Gray,
N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, P. Tillet, F. P. Such, D. Cummings,
M.Plappert,F.Chantzis,E.Barnes,A.Herbert-Voss,W.H.Guss,A.Nichol,A.Paino,N.Tezak,J.Tang,
I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, J. Achiam, V. Misra,
E. Morikawa, A. Radford, M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew,
D. Amodei, S. McCandlish, I. Sutskever, and W. Zaremba. Evaluating large language models trained
on code. CoRR, abs/2107.03374, 2021. URL https://arxiv.org/abs/2107.03374 .
A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung,
C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint
arXiv:2204.02311 , 2022.
T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for
transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Min-
derer,G.Heigold,S.Gelly,etal. Animageisworth16x16words: Transformersforimagerecognition
at scale.arXiv preprint arXiv:2010.11929 , 2020.
T.Ge,H.Xia,X.Sun,S.Chen,andF.Wei. Losslessaccelerationforseq2seqgenerationwithaggressive
decoding. ArXiv, abs/2205.10350, 2022.
8
Accelerating Large Language Model Decoding with Speculative Sampling
J. Hoﬀmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A.
Hendricks,J.Welbl,A.Clark,etal. Trainingcompute-optimallargelanguagemodels. arXivpreprint
arXiv:2203.15556 , 2022.
X. Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu. TinyBERT: Distilling BERT
for natural language understanding. In Findings of the Association for Computational Linguis-
tics: EMNLP 2020 , pages 4163–4174, Online, Nov. 2020. Association for Computational Linguis-
tics. doi: 10.18653/v1/2020.ﬁndings-emnlp.372. URL https://aclanthology.org/2020.
findings-emnlp.372 .
Y. Kim and A. M. Rush. Sequence-level knowledge distillation. CoRR, abs/1606.07947, 2016. URL
http://arxiv.org/abs/1606.07947 .
Y. Leviathan, M. Kalman, and Y. Matias. Fast inference from transformers via speculative decoding.
ArXiv, abs/2211.17192, 2022.
Y. Levine, N. Wies, O. Sharir, H. Bata, and A. Shashua. The depth-to-width interplay in self-attention.
arXiv preprint arXiv:2006.12467 , 2020.
S. Narayan, S. B. Cohen, and M. Lapata. Don’t give me the details, just the summary! topic-
aware convolutional neural networks for extreme summarization. In Proceedings of the 2018
Conference on Empirical Methods in Natural Language Processing , pages 1797–1807, Brussels,
Belgium, Oct.-Nov. 2018. Association for Computational Linguistics. doi: 10.18653/v1/D18-1206.
URLhttps://aclanthology.org/D18-1206 .
R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal,
and J. Dean. Eﬃciently scaling transformer inference. arXiv preprint arXiv:2211.05102 , 2022.
J. W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoﬀmann, F. Song, J. Aslanides, S. Henderson, R. Ring,
S. Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv
preprint arXiv:2112.11446 , 2021.
V. Sanh, L. Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster,
cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.
N. Shazeer. Fast transformer decoding: One write-head is all you need. CoRR, abs/1911.02150, 2019.
URLhttp://arxiv.org/abs/1911.02150 .
M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-lm: Training
multi-billionparameterlanguagemodelsusingmodelparallelism. arXivpreprintarXiv:1909.08053 ,
2019.
Y.Song, C.Meng, R.Liao, andS.Ermon. Acceleratingfeedforwardcomputationviaparallelnonlinear
equationsolving. InM.MeilaandT.Zhang,editors, Proceedingsofthe38thInternationalConference
on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 9791–9800.
PMLR, 18–24 Jul 2021. URL https://proceedings.mlr.press/v139/song21a.html .
M. Stern, N. Shazeer, and J. Uszkoreit. Blockwise parallel decoding for deep autoregressive models.
CoRR, abs/1811.03115, 2018. URL http://arxiv.org/abs/1811.03115 .
A. Wiggers and E. Hoogeboom. Predictive sampling with forecasting autoregressive models. In H. D.
III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning ,
volume 119 of Proceedings of Machine Learning Research , pages 10260–10269. PMLR, 13–18 Jul
2020. URL https://proceedings.mlr.press/v119/wiggers20a.html .
9
Accelerating Large Language Model Decoding with Speculative Sampling
Z. Yao, R. Y. Aminabadi, M. Zhang, X. Wu, C. Li, and Y. He. Zeroquant: Eﬃcient and aﬀordable
post-training quantization for large-scale transformers. arXiv preprint arXiv:2206.01861 , 2022.
Supplementary Materials
Author Contributions
•Initial proposal: Charlie Chen, John Jumper and Geoﬀrey Irving
•Initial Implementation, Optimisation and Scaling: Charlie Chen
•Modiﬁed Rejection Sampling Scheme: John Jumper
•Engineering Improvements: Jean-Baptiste Lespiau and Charlie Chen
•Experiments: Charlie Chen, Sebastian Borgeaud and Laurent Sifre
•Draft of Manuscript: Charlie Chen and Sebastian Borgeaud
•Manuscript Feedback: Laurent Sifre, Geoﬀrey Irving and John Jumper
Acknowledgements
We’d like to thank Oriol Vinyals and Koray Kavukcuoglu for your kind advice and leadership. We’d
also like to thank Evan Senter for your additional feedback on the manuscript and Amelia Glaese for
your support in navigating the publishing process. Finally, we’d like to thank Blake Hechtman, Berkin
Ilbeyi for your valuable advice on XLA and Nikolai Grigoriev for our discussions on the various tricks
that can be applied to the transformer architecture.
Hyperparams
Table 2jHyperparameters for the draft model
Model 𝑑modelHeads Layers Params
Target (Chinchilla) 8192 64 80 70B
Draft 6144 48 8 4B
Proofs
Theorem1 (ModiﬁedRejectionSamplingrecoversthetargetdistribution) .Givendiscretedistributions
𝑞,𝑝and a single draft sample ˜𝑥𝑝, let𝑋be the ﬁnal resulting sample. For 𝑋=𝑥to be true, we must
either sample ˜𝑥=𝑥and then accept it, or resample it after ˜𝑥(of any value) is rejected. Hence:
ℙ¹𝑋=𝑥º
=ℙ¹˜𝑥=𝑥ºℙ¹˜𝑥acceptedj˜𝑥=𝑥º¸ℙ¹˜𝑥rejectedºℙ¹𝑋=𝑥j˜𝑥rejectedº
For the ﬁrst term, we apply the acceptance rule:
ℙ¹˜𝑥=𝑥ºℙ¹˜𝑥acceptedj˜𝑥=𝑥º
=𝑝¹𝑥ºmin
1𝑞¹𝑥º
𝑝¹𝑥º
10
Accelerating Large Language Model Decoding with Speculative Sampling
=min¹𝑝¹𝑥º𝑞¹𝑥ºº
For the second conditional term, we apply the resampling rule:
ℙ¹𝑋=𝑥j˜𝑥rejectedº=¹𝑞¹𝑥º𝑝¹𝑥ºº¸
Where¹º¸denotes:
¹𝑓¹𝑥ºº¸=max¹0 𝑓¹𝑥ººÍ
𝑥max¹0 𝑓¹𝑥ºº
Finally, we calculate the probability of rejection:
ℙ¹˜𝑥rejectedº=1ℙ¹˜𝑥acceptedº
=1∑︁
𝑥0ℙ¹𝑋=𝑥0˜𝑥acceptedº
=1∑︁
𝑥0min¹𝑝¹𝑥0º𝑞¹𝑥0ºº
=∑︁
𝑥0max¹0𝑞¹𝑥0º𝑝¹𝑥0ºº
=∑︁
𝑥0𝑞¹𝑥0ºmin¹𝑝¹𝑥0º𝑞¹𝑥0ºº
=∑︁
𝑥0max¹0𝑞¹𝑥0º𝑝¹𝑥0ºº
This is equal to the denominator of ¹𝑞¹𝑥º𝑝¹𝑥ºº¸, so:
ℙ¹˜𝑥rejectedºℙ¹𝑋=𝑥j˜𝑥rejectedº=max¹0𝑞¹𝑥º𝑝¹𝑥ºº
Hence:
ℙ¹𝑋=𝑥º
=min¹𝑝¹𝑥º𝑞¹𝑥ºº¸max¹0𝑞¹𝑥º𝑝¹𝑥ºº
=𝑞¹𝑥º
and we have recovered the desired target.
11