ViViT: A Video Vision Transformer

Summary:
We present pure-transformer based models for video classification, drawing
upon the recent success of such models in image classification. Our model
extracts spatio-temporal tokens from the input video, which are then encoded by
a series of transformer layers. In order to handle the long sequences of tokens
encountered in video, we propose several, efficient variants of our model which
factorise the spatial- and temporal-dimensions of the input. Although
transformer-based models are known to only be effective when large training
datasets are available, we show how we can effectively regularise the model
during training and leverage pretrained image models to be able to train on
comparatively small datasets. We conduct thorough ablation studies, and achieve
state-of-the-art results on multiple video classification benchmarks including
Kinetics 400 and 600, Epic Kitchens, Something-Something v2 and Moments in
Time, outperforming prior methods based on deep 3D convolutional networks. To
facilitate further research, we release code at
https://github.com/google-research/scenic/tree/main/scenic/projects/vivit

Anurag Arnab
Mostafa Dehghani
Georg Heigold
Chen Sun
Mario Lučić
Cordelia Schmid

ICCV 2021. Code at
  https://github.com/google-research/scenic/tree/main/scenic/projects/vivit

ViViT: A Video Vision Transformer
Anurag Arnab*Mostafa Dehghani*Georg Heigold Chen Sun Mario Lu ˇci´c†Cordelia Schmid†
Google Research
faarnab, dehghani, heigold, chensun, lucic, cordelias g@google.com
Abstract
We present pure-transformer based models for video
classiﬁcation, drawing upon the recent success of such mod-
els in image classiﬁcation. Our model extracts spatio-
temporal tokens from the input video, which are then en-
coded by a series of transformer layers. In order to han-
dle the long sequences of tokens encountered in video, we
propose several, efﬁcient variants of our model which fac-
torise the spatial- and temporal-dimensions of the input. Al-
though transformer-based models are known to only be ef-
fective when large training datasets are available, we show
how we can effectively regularise the model during training
and leverage pretrained image models to be able to train on
comparatively small datasets. We conduct thorough abla-
tion studies, and achieve state-of-the-art results on multiple
video classiﬁcation benchmarks including Kinetics 400 and
600, Epic Kitchens, Something-Something v2 and Moments
in Time, outperforming prior methods based on deep 3D
convolutional networks. To facilitate further research, we
release code at https://github.com/google-research/scenic.
1. Introduction
Approaches based on deep convolutional neural net-
works have advanced the state-of-the-art across many stan-
dard datasets for vision problems since AlexNet [38]. At
the same time, the most prominent architecture of choice in
sequence-to-sequence modelling (e.g. in natural language
processing) is the transformer [68], which does not use con-
volutions, but is based on multi-headed self-attention. This
operation is particularly effective at modelling long-range
dependencies and allows the model to attend over all ele-
ments in the input sequence. This is in stark contrast to
convolutions where the corresponding “receptive ﬁeld” is
limited, and grows linearly with the depth of the network.
The success of attention-based models in NLP has re-
cently inspired approaches in computer vision to integrate
transformers into CNNs [75, 7], as well as some attempts to
replace convolutions completely [49, 3, 53]. However, it is
*Equal contribution
†Equal advisingonly very recently with the Vision Transformer (ViT) [18],
that a pure-transformer based architecture has outperformed
its convolutional counterparts in image classiﬁcation. Doso-
vitskiy et al. [18] closely followed the original transformer
architecture of [68], and noticed that its main beneﬁts
were observed at large scale – as transformers lack some
of the inductive biases of convolutions (such as transla-
tional equivariance), they seem to require more data [18]
or stronger regularisation [64].
Inspired by ViT, and the fact that attention-based ar-
chitectures are an intuitive choice for modelling long-
range contextual relationships in video, we develop sev-
eral transformer-based models for video classiﬁcation. Cur-
rently, the most performant models are based on deep 3D
convolutional architectures [8, 20, 21] which were a natu-
ral extension of image classiﬁcation CNNs [27, 60]. Re-
cently, these models were augmented by incorporating self-
attention into their later layers to better capture long-range
dependencies [75, 23, 79, 1].
As shown in Fig. 1, we propose pure-transformer mod-
els for video classiﬁcation. The main operation performed
in this architecture is self-attention, and it is computed on
a sequence of spatio-temporal tokens that we extract from
the input video. To effectively process the large number of
spatio-temporal tokens that may be encountered in video,
we present several methods of factorising our model along
spatial and temporal dimensions to increase efﬁciency and
scalability. Furthermore, to train our model effectively on
smaller datasets, we show how to reguliarise our model dur-
ing training and leverage pretrained image models.
We also note that convolutional models have been de-
veloped by the community for several years, and there are
thus many “best practices” associated with such models.
As pure-transformer models present different characteris-
tics, we need to determine the best design choices for such
architectures. We conduct a thorough ablation analysis of
tokenisation strategies, model architecture and regularisa-
tion methods. Informed by this analysis, we achieve state-
of-the-art results on multiple standard video classiﬁcation
benchmarks, including Kinetics 400 and 600 [35], Epic
Kitchens 100 [13], Something-Something v2 [26] and Mo-
ments in Time [45].arXiv:2103.15691v2  [cs.CV]  1 Nov 2021
…1CLS032NPosition + Token EmbeddingMLP HeadClassFactorisedEncoderL×
K         V         Q Self-AttentionTransformer  Encoder MLPLayer Norm
Layer NormMulti-HeadDot-Product Attention
Embed to tokensFactorisedSelf-Attention
21NFactorisedDot-Product
●●●SpatialSpatial●●●TemporalTemporal
SpatialTemporalSpatialTemporal●●●SpatialTemporal●●●FuseSpatialTemporalFuse
●●●21N●●●21N●●●Figure 1: We propose a pure-transformer architecture for video classiﬁcation, inspired by the recent success of such models for images [18].
To effectively process a large number of spatio-temporal tokens, we develop several model variants which factorise different components
of the transformer encoder over the spatial- and temporal-dimensions. As shown on the right, these factorisations correspond to different
attention patterns over space and time.
2. Related Work
Architectures for video understanding have mirrored ad-
vances in image recognition. Early video research used
hand-crafted features to encode appearance and motion
information [41, 69]. The success of AlexNet on Ima-
geNet [38, 16] initially led to the repurposing of 2D im-
age convolutional networks (CNNs) for video as “two-
stream” networks [34, 56, 47]. These models processed
RGB frames and optical ﬂow images independently before
fusing them at the end. Availability of larger video classi-
ﬁcation datasets such as Kinetics [35] subsequently facili-
tated the training of spatio-temporal 3D CNNs [8, 22, 65]
which have signiﬁcantly more parameters and thus require
larger training datasets. As 3D convolutional networks re-
quire signiﬁcantly more computation than their image coun-
terparts, many architectures factorise convolutions across
spatial and temporal dimensions and/or use grouped convo-
lutions [59, 66, 67, 81, 20]. We also leverage factorisation
of the spatial and temporal dimensions of videos to increase
efﬁciency, but in the context of transformer-based models.
Concurrently, in natural language processing (NLP),
Vaswani et al. [68] achieved state-of-the-art results by re-
placing convolutions and recurrent networks with the trans-
former network that consisted only of self-attention, layer
normalisation and multilayer perceptron (MLP) operations.
Current state-of-the-art architectures in NLP [17, 52] re-
main transformer-based, and have been scaled to web-scale
datasets [5]. Many variants of the transformer have also
been proposed to reduce the computational cost of self-
attention when processing longer sequences [10, 11, 37,
62, 63, 73] and to improve parameter efﬁciency [40, 14].
Although self-attention has been employed extensively in
computer vision, it has, in contrast, been typically incor-
porated as a layer at the end or in the later stages of
the network [75, 7, 32, 77, 83] or to augment residualblocks [30, 6, 9, 57] within a ResNet architecture [27].
Although previous works attempted to replace convolu-
tions in vision architectures [49, 53, 55], it is only very re-
cently that Dosovitisky et al. [18] showed with their ViT ar-
chitecture that pure-transformer networks, similar to those
employed in NLP, can achieve state-of-the-art results for
image classiﬁcation too. The authors showed that such
models are only effective at large scale, as transformers lack
some of inductive biases of convolutional networks (such
as translational equivariance), and thus require datasets
larger than the common ImageNet ILSRVC dataset [16] to
train. ViT has inspired a large amount of follow-up work
in the community, and we note that there are a number
of concurrent approaches on extending it to other tasks in
computer vision [71, 74, 84, 85] and improving its data-
efﬁciency [64, 48]. In particular, [4, 46] have also proposed
transformer-based models for video.
In this paper, we develop pure-transformer architectures
for video classiﬁcation. We propose several variants of our
model, including those that are more efﬁcient by factoris-
ing the spatial and temporal dimensions of the input video.
We also show how additional regularisation and pretrained
models can be used to combat the fact that video datasets
are not as large as their image counterparts that ViT was
originally trained on. Furthermore, we outperform the state-
of-the-art across ﬁve popular datasets.
3. Video Vision Transformers
We start by summarising the recently proposed Vision
Transformer [18] in Sec. 3.1, and then discuss two ap-
proaches for extracting tokens from video in Sec. 3.2. Fi-
nally, we develop several transformer-based architectures
for video classiﬁcation in Sec. 3.3 and 3.4.
3.1. Overview of Vision Transformers (ViT)
Vision Transformer (ViT) [18] adapts the transformer
architecture of [68] to process 2D images with minimal
changes. In particular, ViT extracts Nnon-overlapping im-
age patches, xi2Rhw, performs a linear projection and
then rasterises them into 1D tokens zi2Rd. The sequence
of tokens input to the following transformer encoder is
z= [zcls;Ex1;Ex2;:::;ExN] +p; (1)
where the projection by Eis equivalent to a 2D convolution.
As shown in Fig. 1, an optional learned classiﬁcation token
zclsis prepended to this sequence, and its representation at
the ﬁnal layer of the encoder serves as the ﬁnal represen-
tation used by the classiﬁcation layer [17]. In addition, a
learned positional embedding, p2RNd, is added to the
tokens to retain positional information, as the subsequent
self-attention operations in the transformer are permutation
invariant. The tokens are then passed through an encoder
consisting of a sequence of Ltransformer layers. Each layer
`comprises of Multi-Headed Self-Attention [68], layer nor-
malisation (LN) [2], and MLP blocks as follows:
y`=MSA (LN(z`)) +z`(2)
z`+1=MLP(LN(y`)) +y`: (3)
The MLP consists of two linear projections separated by a
GELU non-linearity [28] and the token-dimensionality, d,
remains ﬁxed throughout all layers. Finally, a linear classi-
ﬁer is used to classify the encoded input based on zL
cls2Rd,
if it was prepended to the input, or a global average pooling
of all the tokens, zL, otherwise.
As the transformer [68], which forms the basis of
ViT [18], is a ﬂexible architecture that can operate on any
sequence of input tokens z2RNd, we describe strategies
for tokenising videos next.
3.2. Embedding video clips
We consider two simple methods for mapping a video
V2RTHWCto a sequence of tokens ~ z2
Rntnhnwd. We then add the positional embedding and
reshape into RNdto obtain z, the input to the transformer.
Uniform frame sampling As illustrated in Fig. 2, a
straightforward method of tokenising the input video is to
uniformly sample ntframes from the input video clip, em-
bed each 2D frame independently using the same method
as ViT [18], and concatenate all these tokens together. Con-
cretely, ifnhnwnon-overlapping image patches are ex-
tracted from each frame, as in [18], then a total of ntnhnw
tokens will be forwarded through the transformer encoder.
Intuitively, this process may be seen as simply constructing
a large 2D image to be tokenised following ViT. We note
that this is the input embedding method employed by the
concurrent work of [4].
#!"Figure 2: Uniform frame sampling: We simply sample ntframes,
and embed each 2D frame independently following ViT [18].
!"
#
Figure 3: Tubelet embedding. We extract and linearly embed non-
overlapping tubelets that span the spatio-temporal input volume.
Tubelet embedding An alternate method, as shown in
Fig. 3, is to extract non-overlapping, spatio-temporal
“tubes” from the input volume, and to linearly project this to
Rd. This method is an extension of ViT’s embedding to 3D,
and corresponds to a 3D convolution. For a tubelet of di-
mensionthw,nt=bT
tc,nh=bH
hcandnw=bW
wc,
tokens are extracted from the temporal, height, and width
dimensions respectively. Smaller tubelet dimensions thus
result in more tokens which increases the computation.
Intuitively, this method fuses spatio-temporal information
during tokenisation, in contrast to “Uniform frame sam-
pling” where temporal information from different frames is
fused by the transformer.
3.3. Transformer Models for Video
As illustrated in Fig. 1, we propose multiple transformer-
based architectures. We begin with a straightforward ex-
tension of ViT [18] that models pairwise interactions be-
tween all spatio-temporal tokens, and then develop more
efﬁcient variants which factorise the spatial and temporal
dimensions of the input video at various levels of the trans-
former architecture.
Model 1: Spatio-temporal attention This model sim-
ply forwards all spatio-temporal tokens extracted from the
video, z0, through the transformer encoder. We note that
this has also been explored concurrently by [4] in their
“Joint Space-Time” model. In contrast to CNN architec-
tures, where the receptive ﬁeld grows linearly with the
number of layers, each transformer layer models all pair-
…1CLSNPositional + Token EmbeddingTemporal + Token Embedding
Embed to tokens…1N2…1N…TTemporal Transformer EncoderMLP Head
Class…CLS100CLS0CLS0Spatial Transformer EncoderSpatial Transformer EncoderSpatial Transformer EncoderFigure 4: Factorised encoder (Model 2). This model consists of
two transformer encoders in series: the ﬁrst models interactions
between tokens extracted from the same temporal index to produce
a latent representation per time-index. The second transformer
models interactions between time steps. It thus corresponds to a
“late fusion” of spatial- and temporal information.
wise interactions between all spatio-temporal tokens, and it
thus models long-range interactions across the video from
the ﬁrst layer. However, as it models all pairwise in-
teractions, Multi-Headed Self Attention (MSA) [68] has
quadratic complexity with respect to the number of tokens.
This complexity is pertinent for video, as the number of to-
kens increases linearly with the number of input frames, and
motivates the development of more efﬁcient architectures
next.
Model 2: Factorised encoder As shown in Fig. 4, this
model consists of two separate transformer encoders. The
ﬁrst, spatial encoder, only models interactions between to-
kens extracted from the same temporal index. A representa-
tion for each temporal index, hi2Rd, is obtained after Ls
layers: This is the encoded classiﬁcation token, zLs
clsif it was
prepended to the input (Eq. 1), or a global average pooling
from the tokens output by the spatial encoder, zLs, other-
wise. The frame-level representations, hi, are concatenated
intoH2Rntd, and then forwarded through a temporal
encoder consisting of Lttransformer layers to model in-
teractions between tokens from different temporal indices.
The output token of this encoder is then ﬁnally classiﬁed.
This architecture corresponds to a “late fusion” [34,
56, 72, 46] of temporal information, and the initial spa-
tial encoder is identical to the one used for image classi-
ﬁcation. It is thus analogous to CNN architectures such
as [24, 34, 72, 86] which ﬁrst extract per-frame fea-
tures, and then aggregate them into a ﬁnal representation
before classifying them. Although this model has more
transformer layers than Model 1 (and thus more parame-
ters), it requires fewer ﬂoating point operations (FLOPs),
as the two separate transformer blocks have a complexity
ofO((nhnw)2+n2
t)compared toO((ntnhnw)2)of
Model 1.
Transformer Block xLK         V       Q MLPLayer Norm
Layer NormMulti-HeadAttention
K         V       Q Layer NormMulti-HeadAttention
Temporal Self-Attention BlockSpatial Self-Attention BlockToken embeddingPositional embeddingFigure 5: Factorised self-attention (Model 3). Within each trans-
former block, the multi-headed self-attention operation is fac-
torised into two operations (indicated by striped boxes) that ﬁrst
only compute self-attention spatially, and then temporally.
Model 3: Factorised self-attention This model, in con-
trast, contains the same number of transformer layers as
Model 1. However, instead of computing multi-headed
self-attention across all pairs of tokens, z`, at layerl, we
factorise the operation to ﬁrst only compute self-attention
spatially (among all tokens extracted from the same tem-
poral index), and then temporally (among all tokens ex-
tracted from the same spatial index) as shown in Fig. 5.
Each self-attention block in the transformer thus models
spatio-temporal interactions, but does so more efﬁciently
than Model 1 by factorising the operation over two smaller
sets of elements, thus achieving the same computational
complexity as Model 2. We note that factorising attention
over input dimensions has also been explored in [29, 78],
and concurrently in the context of video by [4] in their “Di-
vided Space-Time” model.
This operation can be performed efﬁciently by reshaping
the tokens zfromR1ntnhnwdtoRntnhnwd(denoted
byzs) to compute spatial self-attention. Similarly, the input
to temporal self-attention, ztis reshaped to Rnhnwntd.
Here we assume the leading dimension is the “batch dimen-
sion”. Our factorised self-attention is deﬁned as
y`
s= MSA(LN( z`
s)) +z`
s (4)
y`
t= MSA(LN( y`
s)) +y`
s (5)
z`+1= MLP(LN( y`
t)) +y`
t: (6)
We observed that the order of spatial-then-temporal self-
attention or temporal-then-spatial self-attention does not
make a difference, provided that the model parameters are
initialised as described in Sec. 3.4. Note that the number
of parameters, however, increases compared to Model 1, as
there is an additional self-attention layer (cf. Eq. 7). We do
not use a classiﬁcation token in this model, to avoid ambi-
guities when reshaping the input tokens between spatial and
temporal dimensions.
Model 4: Factorised dot-product attention Finally, we
develop a model which has the same computational com-
plexity as Models 2 and 3, while retaining the same number
of parameters as the unfactorised Model 1. The factorisa-
tion of spatial- and temporal dimensions is similar in spirit
K         V       Q Self-Attention BlockLayer NormMulti-HeadDot-product AttentionConcatenateLinear
K             V                Q Scaled Dot-Product AttentionLinearLinearLinearSpatial HeadsK             V                Q Scaled Dot-Product AttentionLinearLinearLinearTemporal HeadsFigure 6: Factorised dot-product attention (Model 4). For half of
the heads, we compute dot-product attention over only the spatial
axes, and for the other half, over only the temporal axis.
to Model 3, but we factorise the multi-head dot-product at-
tention operation instead (Fig. 6). Concretely, we compute
attention weights for each token separately over the spatial-
and temporal-dimensions using different heads. First, we
note that the attention operation for each head is deﬁned as
Attention( Q;K;V) = SoftmaxQK>
pdk
V: (7)
In self-attention, the queries Q=XW q, keys K=XW k,
and values V=XW vare linear projections of the input X
withX;Q;K;V2RNd. Note that in the unfactorised
case (Model 1), the spatial and temporal dimensions are
merged asN=ntnhnw.
The main idea here is to modify the keys and values for
each query to only attend over tokens from the same spatial-
and temporal index by constructing Ks;Vs2Rnhnwd
andKt;Vt2Rntd, namely the keys and values corre-
sponding to these dimensions. Then, for half of the atten-
tion heads, we attend over tokens from the spatial dimen-
sion by computing Ys= Attention( Q;Ks;Vs), and for
the rest we attend over the temporal dimension by comput-
ingYt= Attention( Q;Kt;Vt). Given that we are only
changing the attention neighbourhood for each query, the
attention operation has the same dimension as in the unfac-
torised case, namely Ys;Yt2RNd. We then combine
the outputs of multiple heads by concatenating them and
using a linear projection [68], Y= Concat( Ys;Yt)WO.
3.4. Initialisation by leveraging pretrained models
ViT [18] has been shown to only be effective when
trained on large-scale datasets, as transformers lack some of
the inductive biases of convolutional networks [18]. How-
ever, even the largest video datasets such as Kinetics [35],
have several orders of magnitude less labelled examples
when compared to their image counterparts [16, 39, 58]. As
a result, training large models from scratch to high accuracy
is extremely challenging. To sidestep this issue, and enable
more efﬁcient training we initialise our video models from
pretrained image models. However, this raises several prac-
tical questions, speciﬁcally on how to initialise parameters
not present or incompatible with image models. We nowdiscuss several effective strategies to initialise these large-
scale video classiﬁcation models.
Positional embeddings A positional embedding pis
added to each input token (Eq. 1). However, our video
models have nttimes more tokens than the pretrained im-
age model. As a result, we initialise the positional embed-
dings by “repeating” them temporally from Rnwnhdto
Rntnhnwd. Therefore, at initialisation, all tokens with
the same spatial index have the same embedding which is
then ﬁne-tuned.
Embedding weights, EWhen using the “tubelet embed-
ding” tokenisation method (Sec. 3.2), the embedding ﬁlter
Eis a 3D tensor, compared to the 2D tensor in the pre-
trained model, Eimage. A common approach for initialising
3D convolutional ﬁlters from 2D ﬁlters for video classiﬁca-
tion is to “inﬂate” them by replicating the ﬁlters along the
temporal dimension and averaging them [8, 22] as
E=1
t[Eimage;:::;Eimage;:::;Eimage]: (8)
We consider an additional strategy, which we denote as
“central frame initialisation”, where Eis initialised with ze-
roes along all temporal positions, except at the centre bt
2c,
E= [0;:::;Eimage;:::;0]: (9)
Therefore, the 3D convolutional ﬁlter effectively behaves
like “Uniform frame sampling” (Sec. 3.2) at initialisation,
while also enabling the model to learn to aggregate temporal
information from multiple frames as training progresses.
Transformer weights for Model 3 The transformer
block in Model 3 (Fig. 5) differs from the pretrained ViT
model [18], in that it contains two multi-headed self atten-
tion (MSA) modules. In this case, we initialise the spatial
MSA module from the pretrained module, and initialise all
weights of the temporal MSA with zeroes, such that Eq. 5
behaves as a residual connection [27] at initialisation.
4. Empirical evaluation
We ﬁrst present our experimental setup and implementa-
tion details in Sec. 4.1, before ablating various components
of our model in Sec. 4.2. We then present state-of-the-art
results on ﬁve datasets in Sec. 4.3.
4.1. Experimental Setup
Network architecture and training Our backbone archi-
tecture follows that of ViT [18] and BERT [17]. We con-
sider ViT-Base (ViT-B, L=12,NH=12,d=768), ViT-Large
(ViT-L,L=24,NH=16,d=1024 ), and ViT-Huge (ViT-H,
L=32,NH=16,d=1280 ), whereLis the number of trans-
former layers, each with a self-attention block of NHheads
Table 1: Comparison of input encoding methods using ViViT-B
and spatio-temporal attention on Kinetics. Further details in text.
Top-1 accuracy
Uniform frame sampling 78.5
Tubelet embedding
Random initialisation [25] 73.2
Filter inﬂation [8] 77.6
Central frame 79.2
and hidden dimension d. We also apply the same naming
scheme to our models (e.g., ViViT-B/16x2 denotes a ViT-
Base backbone with a tubelet size of hwt= 16162).
In all experiments, the tubelet height and width are equal.
Note that smaller tubelet sizes correspond to more tokens at
the input, and thus more computation.
We train our models using synchronous SGD and mo-
mentum, a cosine learning rate schedule and TPU-v3 ac-
celerators. We initialise our models from a ViT image
model trained either on ImageNet-21K [16] (unless other-
wise speciﬁed) or the larger JFT [58] dataset. We imple-
ment our method using the Scenic library [15] and have re-
leased our code and models.
Datasets We evaluate the performance of our proposed
models on a diverse set of video classiﬁcation datasets:
Kinetics [35] consists of 10-second videos sampled at
25fps from YouTube. We evaluate on both Kinetics 400
and 600, containing 400 and 600 classes respectively. As
these are dynamic datasets (videos may be removed from
YouTube), we note our dataset sizes are approximately 267
000 and 446 000 respectively.
Epic Kitchens-100 consists of egocentric videos captur-
ing daily kitchen activities spanning 100 hours and 90 000
clips [13]. We report results following the standard “action
recognition” protocol. Here, each video is labelled with a
“verb” and a “noun” and we therefore predict both cate-
gories using a single network with two “heads”. The top-
scoring verb and action pair predicted by the network form
an “action”, and action accuracy is the primary metric.
Moments in Time [45] consists of 800 000, 3-second
YouTube clips that capture the gist of a dynamic scene in-
volving animals, objects, people, or natural phenomena.
Something-Something v2 (SSv2) [26] contains 220 000
videos, with durations ranging from 2 to 6 seconds. In con-
trast to the other datasets, the objects and backgrounds in
the videos are consistent across different action classes, and
this dataset thus places more emphasis on a model’s ability
to recognise ﬁne-grained motion cues.
Inference The input to our network is a video clip of 32
frames using a stride of 2, unless otherwise mentioned, sim-
ilar to [21, 20]. Following common practice, at inference
time, we process multiple views of a longer video and aver-Table 2: Comparison of model architectures using ViViT-B as the
backbone, and tubelet size of 162. We report Top-1 accuracy on
Kinetics 400 (K400) and action accuracy on Epic Kitchens (EK).
Runtime is during inference on a TPU-v3.
K400 EKFLOPs
(109)Params
(106)Runtime
(ms)
Model 1: Spatio-temporal 80.0 43.1 455.2 88.9 58.9
Model 2: Fact. encoder 78.8 43.7 284.4 115.1 17.4
Model 3: Fact. self-attention 77.4 39.1 372.3 117.3 31.7
Model 4: Fact. dot product 76.3 39.5 277.1 88.9 22.9
Model 2: Ave. pool baseline 75.8 38.8 283.9 86.7 17.3
Table 3: The effect of varying the number of temporal transform-
ers,Lt, in the Factorised encoder model (Model 2). We report the
Top-1 accuracy on Kinetics 400. Note that Lt= 0corresponds to
the “average pooling baseline”.
Lt 0 1 4 8 12
Top-1 75.8 78.6 78.8 78.8 78.9
age per-view logits to obtain the ﬁnal result. Unless other-
wise speciﬁed, we use a total of 4 views per video (as this
is sufﬁcient to “see” the entire video clip across the various
datasets), and ablate these and other design choices next.
4.2. Ablation study
Input encoding We ﬁrst consider the effect of different
input encoding methods (Sec. 3.2) using our unfactorised
model (Model 1) and ViViT-B on Kinetics 400. As we pass
32-frame inputs to the network, sampling 8 frames and ex-
tracting tubelets of length t= 4 correspond to the same
number of tokens in both cases. Table 1 shows that tubelet
embedding initialised using the “central frame” method
(Eq. 9) performs well, outperforming the commonly-used
“ﬁlter inﬂation” initialisation method [8, 22] by 1.6%, and
“uniform frame sampling” by 0.7%. We therefore use this
encoding method for all subsequent experiments.
Model variants We compare our proposed model vari-
ants (Sec. 3.3) across the Kinetics 400 and Epic Kitchens
datasets, both in terms of accuracy and efﬁciency, in Tab. 2.
In all cases, we use the “Base” backbone and tubelet size of
162. Model 2 (“Factorised Encoder”) has an additional
hyperparameter, the number of temporal transformers, Lt.
We setLt= 4 for all experiments and show in Tab. 3 that
the model is not sensitive to this choice.
The unfactorised model (Model 1) performs the best
on Kinetics 400. However, it can also overﬁt on smaller
datasets such as Epic Kitchens, where we ﬁnd our “Fac-
torised Encoder” (Model 2) to perform the best. We also
consider an additional baseline (last row), based on Model
2, where we do not use any temporal transformer, and sim-
ply average pool the frame-level representations from the
spatial encoder before classifying. This average pooling
baseline performs the worst, and has a larger accuracy drop
Table 4: The effect of progressively adding regularisation (each
row includes all methods above it) on Top-1 action accuracy on
Epic Kitchens. We use a Factorised encoder model with tubelet
size162.
Top-1 accuracy
Random crop, ﬂip, colour jitter 38.4
+ Kinetics 400 initialisation 39.6
+ Stochastic depth [31] 40.2
+ Random augment [12] 41.1
+ Label smoothing [61] 43.1
+ Mixup [82] 43.7
16x8 16x4 16x2
Input tubelet size787980Top-1 Accuracy
16x8 16x4 16x2
Input tubelet size0.51.01.5TFLOPs
ViViT-B ViViT-L
(a) Accuracy (b) Compute
Figure 7: The effect of the backbone architecture on (a) accuracy
and (b) computation on Kinetics 400, for the spatio-temporal at-
tention model (Model 1).
16x8 16x4 16x2
Input tubelet size72.575.077.580.0Top-1 Accuracy
16x8 16x4 16x2
Input tubelet size0.20.4TFLOPs
Spatio-temporal Factorised encoder Factorised self-attention Factorised dot-product
(a) Accuracy (b) Compute
Figure 8: The effect of varying the number of temporal tokens on
(a) accuracy and (b) computation on Kinetics 400, for different
variants of our model with a ViViT-B backbone.
on Epic Kitchens, suggesting that this dataset requires more
detailed modelling of temporal relations.
As described in Sec. 3.3, all factorised variants of our
model use signiﬁcantly fewer FLOPs than the unfactorised
Model 1, as the attention is computed separately over
spatial- and temporal-dimensions. Model 4 adds no addi-
tional parameters to the unfactorised Model 1, and uses the
least compute. The temporal transformer encoder in Model
2 operates on only nttokens, which is why there is a barely
a change in compute and runtime over the average pool-
ing baseline, even though it improves the accuracy substan-
tially (3% on Kinetics and 4.9% on Epic Kitchens). Fi-
nally, Model 3 requires more compute and parameters than
the other factorised models, as its additional self-attention
block means that it performs another query-, key-, value-
and output-projection in each transformer layer [68].
Model regularisation Pure-transformer architectures
such as ViT [18] are known to require large training
datasets, and we observed overﬁtting on smaller datasets
like Epic Kitchens and SSv2, even when using an ImageNet
pretrained model. In order to effectively train our modelsTable 5: The effect of spatial resolution on the performance of
ViViT-L/16x2 and spatio-temporal attention on Kinetics 400.
Crop size 224 288 320
Accuracy 80.3 80.7 81.0
GFLOPs 1446 2919 3992
Runtime 58.9 147.6 238.8
on such datasets, we employed several regularisation strate-
gies that we ablate using our “Factorised encoder” model
in Tab. 4. We note that these regularisers were originally
proposed for training CNNs, and that [64] have recently
explored them for training ViT for image classiﬁcation.
Each row of Tab. 4 includes all the methods from the
rows above it, and we observe progressive improvements
from adding each regulariser. Overall, we obtain a substan-
tial overall improvement of 5.3% on Epic Kitchens. We
also achieve a similar improvement of 5% on SSv2 by us-
ing all the regularisation in Tab. 4. Note that the Kinetics-
pretrained models that we initialise from are from Tab. 2,
and that all Epic Kitchens models in Tab. 2 were trained
with all the regularisers in Tab. 4. For larger datasets like
Kinetics and Moments in Time, we do not use these ad-
ditional regularisers (we use only the ﬁrst row of Tab. 4),
as we obtain state-of-the-art results without them. The ap-
pendix contains hyperparameter values and additional de-
tails for all regularisers.
Varying the backbone Figure 7 compares the ViViT-
B and ViViT-L backbones for the unfactorised spatio-
temporal model. We observe consistent improvements in
accuracy as the backbone capacity increases. As expected,
the compute also grows as a function of the backbone size.
Varying the number of tokens We ﬁrst analyse the per-
formance as a function of the number of tokens along the
temporal dimension in Fig. 8. We observe that using smaller
input tubelet sizes (and therefore more tokens) leads to con-
sistent accuracy improvements across all of our model ar-
chitectures. At the same time, computation in terms of
FLOPs increases accordingly, and the unfactorised model
(Model 1) is impacted the most.
We then vary the number of tokens fed into the model
by increasing the spatial crop-size from the default of 224
to 320 in Tab. 5. As expected, there is a consistent increase
in both accuracy and computation. We note that when com-
paring to prior work we consistently obtain state-of-the-art
results (Sec. 4.3) using a spatial resolution of 224, but we
also highlight that further improvements can be obtained at
higher spatial resolutions.
Varying the number of input frames In our experiments
so far, we have kept the number of input frames ﬁxed at 32.
We now increase the number of frames input to the model,
thereby increasing the number of tokens proportionally.
Table 6: Comparisons to state-of-the-art across multiple datasets. For “views”, xydenotes xtemporal crops and yspatial crops. We
report the TFLOPs to process all spatio-temporal views. “FE” denotes our Factorised Encoder model.
(a) Kinetics 400
Method Top 1 Top 5 Views TFLOPs
blVNet [19] 73.5 91.2 – –
STM [33] 73.7 91.6 – –
TEA [42] 76.1 92.5 103 2.10
TSM-ResNeXt-101 [43] 76.3 – – –
I3D NL [75] 77.7 93.3 103 10.77
CorrNet-101 [70] 79.2 – 103 6.72
ip-CSN-152 [66] 79.2 93.8 103 3.27
LGD-3D R101 [51] 79.4 94.4 – –
SlowFast R101-NL [21] 79.8 93.9 103 7.02
X3D-XXL [20] 80.4 94.6 103 5.82
TimeSformer-L [4] 80.7 94.7 13 7.14
ViViT-L/16x2 FE 80.6 92.7 11 3.98
ViViT-L/16x2 FE 81.7 93.8 13 11.94
Methods with large-scale pretraining
ip-CSN-152 [66] (IG [44]) 82.5 95.3 103 3.27
ViViT-L/16x2 FE (JFT) 83.5 94.3 13 11.94
ViViT-H/14x2 (JFT) 84.9 95.8 43 47.77(b) Kinetics 600
Method Top 1 Top 5
AttentionNAS [76] 79.8 94.4
LGD-3D R101 [51] 81.5 95.6
SlowFast R101-NL [21] 81.8 95.1
X3D-XL [20] 81.9 95.5
TimeSformer-L [4] 82.2 95.6
ViViT-L/16x2 FE 82.9 94.6
ViViT-L/16x2 FE (JFT) 84.3 94.9
ViViT-H/14x2 (JFT) 85.8 96.5
(c) Moments in Time
Top 1 Top 5
TSN [72] 25.3 50.1
TRN [86] 28.3 53.4
I3D [8] 29.5 56.1
blVNet [19] 31.4 59.3
AssembleNet-101 [54] 34.3 62.7
ViViT-L/16x2 FE 38.5 64.1(d) Epic Kitchens 100 Top 1 accuracy
Method Action Verb Noun
TSN [72] 33.2 60.2 46.0
TRN [86] 35.3 65.9 45.4
TBN [36] 36.7 66.0 47.2
TSM [43] 38.3 67.9 49.0
SlowFast [21] 38.5 65.6 50.0
ViViT-L/16x2 FE 44.0 66.4 56.8
(e) Something-Something v2
Method Top 1 Top 5
TRN [86] 48.8 77.6
SlowFast [20, 80] 61.7 –
TimeSformer-HR [4] 62.5 –
TSM [43] 63.4 88.5
STM [33] 64.2 89.8
TEA [42] 65.1 –
blVNet [19] 65.2 90.3
ViVIT-L/16x2 FE 65.9 89.9
1 2 3 4 5 6 7
Number of views767880Top-1 Accuracy
32 stride 2 64 stride 2 128 stride 2
Figure 9: The effect of varying the number of frames input to the
network and increasing the number of tokens proportionally. We
use ViViT-L/16x2 Factorised Encoder on Kinetics 400. A Kinetics
video contains 250 frames (10 seconds sampled at 25 fps) and the
accuracy for each model saturates once the number of equidistant
temporal views is sufﬁcient to “see” the whole video clip. Ob-
serve how models processing more frames (and thus more tokens)
achieve higher single- and multi-view accuracy.
Figure 9 shows that as we increase the number of frames
input to the network, the accuracy from processing a sin-
gle view increases, since the network incorporates longer
temporal context. However, common practice on datasets
such as Kinetics [21, 75, 42] is to average results over mul-
tiple, shorter “views” of the same video clip. Figure 9 also
shows that the accuracy saturates once the number of views
is sufﬁcient to cover the whole video. As a Kinetics video
consists of 250 frames, and we sample frames with a stride
of 2, our model which processes 128 frames requires just a
single view to “see” the whole video and achieve its maxi-
mum accuarcy.
Note that we used ViViT-L/16x2 Factorised Encoder
(Model 2) here. As this model is more efﬁcient it can pro-
cess more tokens, compared to the unfactorised Model 1
which runs out of memory after 48 frames using tubelet
lengtht= 2 and a “Large” backbone. Models processingmore frames (and thus more tokens) consistently achieve
higher single- and multi-view accuracy, in line with our ob-
servations in previous experiments (Tab. 5, Fig. 8). Mo-
roever, observe that by processing more frames (and thus
more tokens) with Model 2, we are able to achieve higher
accuracy than Model 1 (with fewer total FLOPs as well).
Finally, we observed that for Model 2, the number of
FLOPs effectively increases linearly with the number of in-
put frames as the overall computation is dominated by the
initial Spatial Transformer. As a result, the total number
of FLOPs for the number of temporal views required to
achieve maximum accuracy is constant across the models.
In other words, ViViT-L/16x2 FE with 32 frames requires
995.3 GFLOPs per view, and 4 views to saturate multi-view
accuracy. The 128-frame model requires 3980.4 GFLOPs
but only a single view. As shown by Fig. 9, the latter model
achieves the highest accuracy.
4.3. Comparison to state-of-the-art
Based on our ablation studies in the previous section,
we compare to the current state-of-the-art using two of our
model variants. We primarily use our Factorised Encoder
model (Model 2), as it can process more tokens than Model
1 to achieve higher accuracy.
Kinetics Tables 6a and 6b show that our spatio-temporal
attention models outperform the state-of-the-art on Kinetics
400 and 600 respectively. Following standard practice, we
take 3 spatial crops (left, centre and right) [21, 20, 66, 75]
for each temporal view, and notably, we require signiﬁ-
cantly fewer views than previous CNN-based methods.
We surpass the previous CNN-based state-of-the-art us-
ing ViViT-L/16x2 Factorised Encoder (FE) pretrained on
ImageNet, and also outperform [4] who concurrently pro-
posed a pure-transformer architecture. Moreover, by initial-
ising our backbones from models pretrained on the larger
JFT dataset [58], we obtain further improvements. Al-
though these models are not directly comparable to previ-
ous work, we do also outperform [66] who pretrained on
the large-scale, Instagram dataset [44]. Our best model uses
a ViViT-H backbone pretrained on JFT and signiﬁcantly ad-
vances the best reported results on Kinetics 400 and 600 to
84.9% and 85.8%, respectively.
Moments in Time We surpass the state-of-the-art by a
signiﬁcant margin as shown in Tab. 6c. We note that the
videos in this dataset are diverse and contain signiﬁcant la-
bel noise, making this task challenging and leading to lower
accuracies than on other datasets.
Epic Kitchens 100 Table 6d shows that our Factorised
Encoder model outperforms previous methods by a signiﬁ-
cant margin. In addition, our model obtains substantial im-
provements for Top-1 accuracy of “noun” classes, and the
only method which achieves higher “verb” accuracy used
optical ﬂow as an additional input modality [43, 50]. Fur-
thermore, all variants of our model presented in Tab. 2 out-
performed the existing state-of-the-art on action accuracy.
We note that we use the same model to predict verbs and
nouns using two separate “heads”, and for simplicity, we do
not use separate loss weights for each head.
Something-Something v2 (SSv2) Finally, Tab. 6e shows
that we achieve state-of-the-art Top-1 accuracy with our
Factorised encoder model (Model 2), albeit with a smaller
margin compared to previous methods. Notably, our Fac-
torised encoder model signiﬁcantly outperforms the concur-
rent TimeSformer [4] method by 2.9%, which also proposes
a pure-transformer model, but does not consider our Fac-
torised encoder variant or our additional regularisation.
SSv2 differs from other datasets in that the backgrounds
and objects are quite similar across different classes, mean-
ing that recognising ﬁne-grained motion patterns is neces-
sary to distinguish classes from each other. Our results sug-
gest that capturing these ﬁne-grained motions is an area of
improvement and future work for our model. We also note
an inverse correlation between the relative performance of
previous methods on SSv2 (Tab. 6e) and Kinetics (Tab. 6a)
suggesting that these two datasets evaluate complementary
characteristics of a model.
5. Conclusion and Future Work
We have presented four pure-transformer models for
video classiﬁcation, with different accuracy and efﬁciency
proﬁles, achieving state-of-the-art results across ﬁve pop-
ular datasets. Furthermore, we have shown how to ef-
fectively regularise such high-capacity models for training
on smaller datasets and thoroughly ablated our main de-
sign choices. Future work is to remove our dependence onimage-pretrained models. Finally, going beyond video clas-
siﬁcation towards more complex tasks is a clear next step.
References
[1] Anurag Arnab, Chen Sun, and Cordelia Schmid. Uniﬁed
graph structured models for video understanding. In ICCV ,
2021. 1
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.
Layer normalization. In arXiv preprint arXiv:1607.06450 ,
2016. 3
[3] Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens,
and Quoc V Le. Attention augmented convolutional net-
works. In ICCV , 2019. 1
[4] Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is
space-time attention all you need for video understanding?
InarXiv preprint arXiv:2102.05095 , 2021. 2, 3, 4, 8, 9
[5] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Sub-
biah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakan-
tan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini
Agarwal, et al. Language models are few-shot learners. In
NeurIPS , 2020. 2
[6] Yue Cao, Jiarui Xu, Stephen Lin, Fangyun Wei, and Han
Hu. Gcnet: Non-local networks meet squeeze-excitation net-
works and beyond. In CVPR Workshops , 2019. 2
[7] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas
Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-
end object detection with transformers. In ECCV , 2020. 1,
2
[8] Joao Carreira and Andrew Zisserman. Quo vadis, action
recognition? a new model and the kinetics dataset. In CVPR ,
2017. 1, 2, 5, 6, 8
[9] Yunpeng Chen, Yannis Kalantidis, Jianshu Li, Shuicheng
Yan, and Jiashi Feng. A2-nets: Double attention networks.
InNeurIPS , 2018. 2
[10] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.
Generating long sequences with sparse transformers. In
arXiv preprint arXiv:1904.10509 , 2019. 2
[11] Krzysztof Choromanski, Valerii Likhosherstov, David Do-
han, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter
Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser,
et al. Rethinking attention with performers. In ICLR , 2021.
2
[12] Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V .
Le. Randaugment: Practical automated data augmentation
with a reduced search space. In NeurIPS , 2020. 7, 13, 14
[13] Dima Damen, Hazel Doughty, Giovanni Maria Farinella,
Antonino Furnari, Jian Ma, Evangelos Kazakos, Davide
Moltisanti, Jonathan Munro, Toby Perrett, Will Price, and
Michael Wray. Rescaling egocentric vision. In arXiv
preprint arXiv:2006.13256 , 2020. 1, 6
[14] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob
Uszkoreit, and Łukasz Kaiser. Universal transformers. In
ICLR , 2019. 2
[15] Mostafa Dehghani, Alexey Gritsenko, Anurag Arnab,
Matthias Minderer, and Yi Tay. Scenic: A JAX library
for computer vision research and beyond. arXiv preprint
arXiv:2110.11403 , 2021. 6
[16] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li,
and Li Fei-Fei. Imagenet: A large-scale hierarchical image
database. In CVPR , 2009. 2, 5, 6
[17] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina
Toutanova. Bert: Pre-training of deep bidirectional trans-
formers for language understanding. In NAACL , 2019. 2, 3,
5
[18] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov,
Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,
Mostafa Dehghani, Matthias Minderer, Georg Heigold, Syl-
vain Gelly, et al. An image is worth 16x16 words: Trans-
formers for image recognition at scale. In ICLR , 2021. 1, 2,
3, 5, 7
[19] Quanfu Fan, Chun-Fu Chen, Hilde Kuehne, Marco Pistoia,
and David Cox. More is less: Learning efﬁcient video repre-
sentations by big-little network and depthwise temporal ag-
gregation. In NeurIPS , 2019. 8
[20] Christoph Feichtenhofer. X3d: Expanding architectures for
efﬁcient video recognition. In CVPR , 2020. 1, 2, 6, 8
[21] Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and
Kaiming He. Slowfast networks for video recognition. In
ICCV , 2019. 1, 6, 8
[22] Christoph Feichtenhofer, Axel Pinz, and Richard Wildes.
Spatiotemporal residual networks for video action recogni-
tion. In NeurIPS , 2016. 2, 5, 6
[23] Rohit Girdhar, Joao Carreira, Carl Doersch, and Andrew Zis-
serman. Video action transformer network. In CVPR , 2019.
1
[24] Rohit Girdhar and Deva Ramanan. Attentional pooling for
action recognition. In NeurIPS , 2017. 4
[25] Xavier Glorot and Yoshua Bengio. Understanding the difﬁ-
culty of training deep feedforward neural networks. In AIS-
TATS , 2010. 6
[26] Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michal-
ski, Joanna Materzynska, Susanne Westphal, Heuna Kim,
Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz
Mueller-Freitag, et al. The” something something” video
database for learning and evaluating visual common sense.
InICCV , 2017. 1, 6
[27] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
Deep residual learning for image recognition. In CVPR ,
2016. 1, 2, 5
[28] Dan Hendrycks and Kevin Gimpel. Gaussian error linear
units (gelus). In arXiv preprint arXiv:1606.08415 , 2016. 3
[29] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim
Salimans. Axial attention in multidimensional transformers.
InarXiv preprint arXiv:1912.12180 , 2019. 4
[30] Jie Hu, Li Shen, and Gang Sun. Squeeze-and-excitation net-
works. In CVPR , 2018. 2
[31] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian
Weinberger. Deep networks with stochastic depth. In ECCV ,
2016. 7, 13, 14
[32] Zilong Huang, Xinggang Wang, Lichao Huang, Chang
Huang, Yunchao Wei, and Wenyu Liu. Ccnet: Criss-cross
attention for semantic segmentation. In ICCV , 2019. 2
[33] Boyuan Jiang, Mengmeng Wang, Weihao Gan, Wei Wu, and
Junjie Yan. Stm: Spatiotemporal and motion encoding for
action recognition. In ICCV , 2019. 8[34] Andrej Karpathy, George Toderici, Sanketh Shetty, Thomas
Leung, Rahul Sukthankar, and Li Fei-Fei. Large-scale video
classiﬁcation with convolutional neural networks. In CVPR ,
2014. 2, 4
[35] Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang,
Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola,
Tim Green, Trevor Back, Paul Natsev, et al. The ki-
netics human action video dataset. In arXiv preprint
arXiv:1705.06950 , 2017. 1, 2, 5, 6
[36] Evangelos Kazakos, Arsha Nagrani, Andrew Zisserman, and
Dima Damen. Epic-fusion: Audio-visual temporal binding
for egocentric action recognition. In ICCV , 2019. 8
[37] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Re-
former: The efﬁcient transformer. In ICLR , 2020. 2
[38] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.
Imagenet classiﬁcation with deep convolutional neural net-
works. In NeurIPS , volume 25, 2012. 1, 2
[39] Alina Kuznetsova, Hassan Rom, Neil Alldrin, Jasper Ui-
jlings, Ivan Krasin, Jordi Pont-Tuset, Shahab Kamali, Stefan
Popov, Matteo Malloci, Tom Duerig, et al. The open im-
ages dataset v4: Uniﬁed image classiﬁcation, object detec-
tion, and visual relationship detection at scale. IJCV , 2020.
5
[40] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin
Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite
bert for self-supervised learning of language representations.
InICLR , 2020. 2
[41] Ivan Laptev. On space-time interest points. IJCV , 64(2-3),
2005. 2
[42] Yan Li, Bin Ji, Xintian Shi, Jianguo Zhang, Bin Kang, and
Limin Wang. Tea: Temporal excitation and aggregation for
action recognition. In CVPR , 2020. 8
[43] Ji Lin, Chuang Gan, and Song Han. Tsm: Temporal shift
module for efﬁcient video understanding. In ICCV , 2019. 8,
9
[44] Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan,
Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe,
and Laurens Van Der Maaten. Exploring the limits of weakly
supervised pretraining. In ECCV , 2018. 8, 9
[45] Mathew Monfort, Alex Andonian, Bolei Zhou, Kandan Ra-
makrishnan, Sarah Adel Bargal, Tom Yan, Lisa Brown,
Quanfu Fan, Dan Gutfreund, Carl V ondrick, et al. Moments
in time dataset: one million videos for event understanding.
PAMI , 42(2):502–508, 2019. 1, 6
[46] Daniel Neimark, Omri Bar, Maya Zohar, and Dotan As-
selmann. Video transformer network. In arXiv preprint
arXiv:2102.00719 , 2021. 2, 4
[47] Joe Yue-Hei Ng, Matthew Hausknecht, Sudheendra Vi-
jayanarasimhan, Oriol Vinyals, Rajat Monga, and George
Toderici. Beyond short snippets: Deep networks for video
classiﬁcation. In CVPR , 2015. 2
[48] Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jian-
fei Cai. Scalable visual transformers with hierarchical pool-
ing. In arXiv preprint arXiv:2103.10619 , 2021. 2
[49] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz
Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Im-
age transformer. In ICML , 2018. 1, 2
[50] Will Price and Dima Damen. An evaluation of action
recognition models on epic-kitchens. In arXiv preprint
arXiv:1908.00867 , 2019. 9
[51] Zhaofan Qiu, Ting Yao, Chong-Wah Ngo, Xinmei Tian, and
Tao Mei. Learning spatio-temporal representation with local
and global diffusion. In CVPR , 2019. 8
[52] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee,
Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and
Peter J Liu. Exploring the limits of transfer learning with a
uniﬁed text-to-text transformer. JMLR , 2020. 2
[53] Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan
Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone
self-attention in vision models. In NeurIPS , 2019. 1, 2
[54] Michael S Ryoo, AJ Piergiovanni, Mingxing Tan, and Anelia
Angelova. Assemblenet: Searching for multi-stream neural
connectivity in video architectures. In ICLR , 2020. 8
[55] Zhuoran Shen, Irwan Bello, Raviteja Vemulapalli, Xuhui Jia,
and Ching-Hui Chen. Global self-attention networks for im-
age recognition. In arXiv preprint arXiv:2010.03019 , 2021.
2
[56] Karen Simonyan and Andrew Zisserman. Two-stream con-
volutional networks for action recognition in videos. In
NeurIPS , 2014. 2, 4
[57] Aravind Srinivas, Tsung-Yi Lin, Niki Parmar, Jonathon
Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck
transformers for visual recognition. In CVPR , 2021. 2
[58] Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhi-
nav Gupta. Revisiting unreasonable effectiveness of data in
deep learning era. In ICCV , 2017. 5, 6, 9
[59] Lin Sun, Kui Jia, Dit-Yan Yeung, and Bertram E Shi. Human
action recognition using factorized spatio-temporal convolu-
tional networks. In ICCV , 2015. 2
[60] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet,
Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent
Vanhoucke, and Andrew Rabinovich. Going deeper with
convolutions. In CVPR , 2015. 1
[61] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon
Shlens, and Zbigniew Wojna. Rethinking the inception ar-
chitecture for computer vision. In CVPR , 2016. 7, 13, 14
[62] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen,
Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebas-
tian Ruder, and Donald Metzler. Long range arena: A
benchmark for efﬁcient transformers. In arXiv preprint
arXiv:2011.04006 , 2020. 2
[63] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Met-
zler. Efﬁcient transformers: A survey. In arXiv preprint
arXiv:2009.06732 , 2020. 2
[64] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco
Massa, Alexandre Sablayrolles, and Herv ´e J´egou. Training
data-efﬁcient image transformers & distillation through at-
tention. In arXiv preprint arXiv:2012.12877 , 2020. 1, 2, 7
[65] Du Tran, Lubomir Bourdev, Rob Fergus, Lorenzo Torresani,
and Manohar Paluri. Learning spatiotemporal features with
3d convolutional networks. In ICCV , 2015. 2
[66] Du Tran, Heng Wang, Lorenzo Torresani, and Matt Feis-
zli. Video classiﬁcation with channel-separated convolu-
tional networks. In ICCV , 2019. 2, 8, 9[67] Du Tran, Heng Wang, Lorenzo Torresani, Jamie Ray, Yann
LeCun, and Manohar Paluri. A closer look at spatiotemporal
convolutions for action recognition. In CVPR , 2018. 2
[68] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszko-
reit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia
Polosukhin. Attention is all you need. In NeurIPS , 2017. 1,
2, 3, 4, 5, 7
[69] Heng Wang, Alexander Kl ¨aser, Cordelia Schmid, and
Cheng-Lin Liu. Dense trajectories and motion boundary de-
scriptors for action recognition. IJCV , 103(1), 2013. 2
[70] Heng Wang, Du Tran, Lorenzo Torresani, and Matt Feiszli.
Video modeling with correlation networks. In CVPR , 2020.
8
[71] Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, and
Liang-Chieh Chen. Max-deeplab: End-to-end panoptic
segmentation with mask transformers. In arXiv preprint
arXiv:2012.00759 , 2020. 2
[72] Limin Wang, Yuanjun Xiong, Zhe Wang, Yu Qiao, Dahua
Lin, Xiaoou Tang, and Luc Van Gool. Temporal segment
networks: Towards good practices for deep action recogni-
tion. In ECCV , 2016. 4, 8
[73] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and
Hao Ma. Linformer: Self-attention with linear complexity.
InarXiv preprint arXiv:2006.04768 , 2020. 2
[74] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao
Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.
Pyramid vision transformer: A versatile backbone for
dense prediction without convolutions. In arXiv preprint
arXiv:2102.12122 , 2021. 2
[75] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaim-
ing He. Non-local neural networks. In CVPR , 2018. 1, 2,
8
[76] Xiaofang Wang, Xuehan Xiong, Maxim Neumann, AJ Pier-
giovanni, Michael S Ryoo, Anelia Angelova, Kris M Kitani,
and Wei Hua. Attentionnas: Spatiotemporal attention cell
search for video classiﬁcation. In ECCV , 2020. 8
[77] Yuqing Wang, Zhaoliang Xu, Xinlong Wang, Chunhua Shen,
Baoshan Cheng, Hao Shen, and Huaxia Xia. End-to-end
video instance segmentation with transformers. In arXiv
preprint arXiv:2011.14503 , 2020. 2
[78] Dirk Weissenborn, Oscar T ¨ackstr ¨om, and Jakob Uszkoreit.
Scaling autoregressive video models. In ICLR , 2020. 4
[79] Chao-Yuan Wu, Christoph Feichtenhofer, Haoqi Fan, Kaim-
ing He, Philipp Krahenbuhl, and Ross Girshick. Long-term
feature banks for detailed video understanding. In CVPR ,
2019. 1
[80] Chao-Yuan Wu, Ross Girshick, Kaiming He, Christoph Fe-
ichtenhofer, and Philipp Krahenbuhl. A multigrid method
for efﬁciently training video models. In CVPR , 2020. 8
[81] Saining Xie, Chen Sun, Jonathan Huang, Zhuowen Tu, and
Kevin Murphy. Rethinking spatiotemporal feature learning:
Speed-accuracy trade-offs in video classiﬁcation. In ECCV ,
2018. 2
[82] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and
David Lopez-Paz. Mixup: Beyond empirical risk minimiza-
tion. In ICLR , 2018. 7, 13, 14
[83] Li Zhang, Dan Xu, Anurag Arnab, and Philip HS Torr. Dy-
namic graph message passing networks. In CVPR , 2020. 2
[84] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, and
Vladlen Koltun. Point transformer. In arXiv preprint
arXiv:2012.09164 , 2020. 2
[85] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu,
Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao
Xiang, Philip HS Torr, et al. Rethinking semantic segmen-
tation from a sequence-to-sequence perspective with trans-
formers. In arXiv preprint arXiv:2012.15840 , 2020. 2
[86] Bolei Zhou, Alex Andonian, Aude Oliva, and Antonio Tor-
ralba. Temporal relational reasoning in videos. In ECCV ,
2018. 4, 8
Appendix
A. Additional experimental details
In this appendix, we provide additional experimental de-
tails. Section A.1 provides additional details about the reg-
ularisers we used and Sec. A.2 details the training hyper-
paramters used for our experiments.
A.1. Further details about regularisers
In this section, we provide additional details and list the
hyperparameters of the additional regularisers that we em-
ployed in Tab. 4. Hyperparameter values for all our experi-
ments are listed in Tab. 7.
Stochastic depth Stochastic depth regularisation was
originally proposed for training very deep residual net-
works [31]. Intuitively, the outputs of a layer, `, are
“dropped out” with probability, pdrop(`)during training, by
setting the output of the layer to be equal to its input.
Following [31], we linearly increase the probability of
dropping a layer according to its depth within the network,
pdrop(`) =`
Lpdrop; (10)
where`is the index of the layer in the network, and Lis the
total number of layers.
Random augment Random augment [12] randomly ap-
plies data augmentation transformations sequentially to an
input example. We follow the public implementation1, but
modify the data augmentation operations to be temporally
consistent throughout the video (in other words, the same
transformation is applied on each frame of the video).
The authors deﬁne two hyperparameters for Random
augment, “number of layers” , the number of augmentation
transformations to apply sequentially to a video and “mag-
nitude”, the strength of the transformation that is shared
across all augmentation operations. Our values for these
parameters are shown in Tab. 7.
Label smoothing Label smoothing was proposed by [61]
originally to regularise training Inception-v3. Concretely,
the label distribution used during training, ~y, is a mixture
of the one-hot ground-truth label, y, and a uniform distribu-
tion,u, to encourage the network to produce less conﬁdent
predictions during training:
~y= (1)y+u: (11)
There is therefore one scalar hyperparamter, 2[0;1].
1https://github.com/tensorflow/models/blob/
master/official/vision/beta/ops/augment.pyMixup Mixup [82] constructs virtual training examples
which are a convex combination of pairs of training exam-
ples and their labels. Concretely, given (xi;yi)and(xj;yj)
wherexidenotes an input vector and yia one-hot input la-
bel, mixup constructs the virtual training example,
~x=xi+ (1)xj
~y=yi+ (1)yj: (12)
2[0;1], and is sampled from a Beta distribution,
Beta(;). Our choice of the hyperparameter is detailed
in Tab. 7.
A.2. Training hyperparameters
Table 7 details the hyperparamters for all of our ex-
periments. We use synchronous SGD with momentum, a
cosine learning rate schedule with linear warmup, and a
batch size of 64 for all experiments. As aforementioned,
we only employed additional regularisation when training
on the smaller Epic Kitchens and Something-Something v2
datasets.
Table 7: Training hyperparamters for experiments in the main paper. “–” indicates that the regularisation method was not used at all. Values
which are constant across all columns are listed once. Datasets are denoted as follows: K400: Kinetics 400. K600: Kinetics 600. MiT:
Moments in Time. EK: Epic Kitchens. SSv2: Something-Something v2.
K400 K600 MiT EK SSv2
Optimisation
Optimiser Synchronous SGD
Momentum 0.9
Batch size 64
Learning rate schedule cosine with linear warmup
Linear warmup epochs 2.5
Base learning rate 0.1 0.1 0.25 0.5 0.5
Epochs 30 30 10 50 35
Data augmentation
Random crop probability 1.0
Random ﬂip probability 0.5
Scale jitter probability 1.0
Maximum scale 1.33
Minimum scale 0.9
Colour jitter probability 0.8 0.8 0.8 – –
Rand augment number of layers [12] – – – 2 2
Rand augment magnitude [12] – – – 15 20
Other regularisation
Stochastic droplayer rate, pdrop[31] – – – 0.2 0.3
Label smoothing [61] – – – 0.2 0.3
Mixup[82] – – – 0.1 0.3

1. Unified Graph Structured Models for Video Understanding

2. Layer Normalization

3. Attention Augmented Convolutional Networks

Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, Quoc V Le

4. Is Space-Time Attention All You Need for Video Understanding?

Gedas Bertasius, Heng Wang, Lorenzo Torresani

5. Language Models are Few-Shot Learners


Summary:
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

Vivit: A Video Vision Transformer

Language Models are Few-Shot Learners