Space-time Mixing Attention for Video Transformer

This paper is on video recognition using Transformers . Very recent attempts in this area have demonstrated promising results in terms of recognition accuracy , yet they have been also shown to induce , in many cases , significant computational overheads due to the additional modelling of the temporal information . In this work , we propose a Video Transformer model the complexity of which scales linearly with the number of frames in the video sequence and hence induces no overhead compared to an image-based Transformer model . To achieve this , our model makes two approximations to the full space-time attention used in Video Transformers : ( a ) It restricts time attention to a local temporal window and capitalizes on the Transformer ’ s depth to obtain full temporal coverage of the video sequence . ( b ) It uses efficient space-time mixing to attend jointly spatial and temporal locations without inducing any additional cost on top of a spatial-only attention model . We also show how to integrate 2 very lightweight mechanisms for global temporal-only attention which provide additional accuracy improvements at minimal computational cost . We demonstrate that our model produces very high recognition accuracy on the most popular video recognition datasets while at the same time being significantly more efficient than other Video Transformer models . Code for our method is made available here . 1 Introduction . Video recognition – in analogy to image recognition – refers to the problem of recognizing events of interest in video sequences such as human activities . Following the tremendous success of Transformers in sequential data , specifically in Natural Language Processing ( NLP ) [ 39 , 5 ] , Vision Transformers were very recently shown to outperform CNNs for image recognition too [ 48 , 13 , 35 ] , signaling a paradigm shift on how visual understanding models should be constructed . In light of this , in this paper , we propose a Video Transformer model as an appealing and promising solution for improving the accuracy of video recognition models . A direct , natural extension of Vision Transformers to the spatio-temporal domain is to perform the self-attention jointly across all S spatial locations and T temporal locations . Full space-time attention though has complexityO ( T 2S2 ) making such a model computationally heavy and , hence , impractical even when compared with the 3D-based convolutional models . As such , our aim is to exploit the temporal information present in video streams while minimizing the computational burden within the Transformer framework for efficient video recognition . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A baseline solution to this problem is to consider spatial-only attention followed by temporal averaging , which has complexity O ( TS2 ) . Similar attempts to reduce the cost of full space-time attention have been recently proposed in [ 3 , 1 ] . These methods have demonstrated promising results in terms of video recognition accuracy , yet they have been also shown to induce , in most of the cases , significant computational overheads compared to the baseline ( spatial-only ) method due to the additional modelling of the temporal information . Our main contribution in this paper is a Video Transformer model that has complexity O ( TS2 ) and , hence , is as efficient as the baseline model , yet , as our results show , it outperforms recently/concurrently proposed work [ 3 , 1 ] in terms of efficiency ( i.e . accuracy/FLOP ) by significant margins . To achieve this our model makes two approximations to the full space-time attention used in Video Transformers : ( a ) It restricts time attention to a local temporal window and capitalizes on the Transformer ’ s depth to obtain full temporal coverage of the video sequence . ( b ) It uses efficient space-time mixing to attend jointly spatial and temporal locations without inducing any additional cost on top of a spatial-only attention model . Fig . 1 shows the proposed approximation to space-time attention . We also show how to integrate two very lightweight mechanisms for global temporal-only attention , which provide additional accuracy improvements at minimal computational cost . We demonstrate that our model is surprisingly effective in terms of capturing long-term dependencies and producing very high recognition accuracy on the most popular video recognition datasets , including Something-Something-v2 [ 17 ] , Kinetics [ 4 ] and Epic Kitchens [ 9 ] , while at the same time being significantly more efficient than other Video Transformer models . 2 Related work . Video recognition : Standard solutions are based on CNNs and can be broadly classified into two categories : 2D- and 3D-based approaches . 2D-based approaches process each frame independently to extract frame-based features which are then aggregated temporally with some sort of temporal modeling ( e.g . temporal averaging ) performed at the end of the network [ 42 , 26 , 27 ] . The works of [ 26 , 27 ] use the “ shift trick ” [ 45 ] to have some temporal modeling at a layer level . 3D-based approaches [ 4 , 16 , 36 ] are considered the current state-of-the-art as they can typically learn stronger temporal models via 3D convolutions . However , they also incur higher computational and memory costs . To alleviate this , a large body of works attempt to improve their efficiency via spatial and/or temporal factorization [ 38 , 37 , 15 ] . CNN vs ViT : Historically , video recognition approaches tend to mimic the architectures used for image classification ( e.g . from AlexNet [ 23 ] to [ 20 ] or from ResNet [ 18 ] and ResNeXt [ 47 ] to [ 16 ] ) . After revolutionizing NLP [ 39 , 32 ] , very recently , Transformer-based architectures showed promising results on large scale image classification too [ 13 ] . While self-attention and attention were previously used in conjunction with CNNs at a layer or block level [ 6 , 50 , 33 ] , the Vision Transformer ( ViT ) of Dosovitskiy et al . [ 13 ] is the first convolution-free , Transformer-based architecture that achieves state-of-the-art on ImageNet [ 11 ] . Video Transformer : Recently/concurrently with our work , vision transformer architectures , derived from [ 13 ] , were used for video recognition [ 3 , 1 ] , too . Because performing full space-time attention is computationally prohibitive ( i.e . O ( T 2S2 ) ) , their main focus is on reducing this via temporal and spatial factorization . In TimeSformer [ 3 ] , the authors propose applying spatial and temporal attention in an alternating manner reducing the complexity to O ( T 2S + TS2 ) . In a similar fashion , ViViT [ 1 ] explores several avenues for space-time factorization . In addition , they also proposed to adapt the patch embedding process from [ 13 ] to 3D ( i.e . video ) data . Our work proposes a completely different approximation to full space-time attention that is also efficient . To this end , we firstly restrict full space-time attention to a local temporal window which is reminiscent of [ 2 ] but applied here to space-time attention and video recognition 1 . Secondly , we define a local joint space-time attention which we show that can be implemented efficiently via the “ shift trick ” [ 45 ] . 3 Method . Video Transformer : We are given a video clip X ∈ RT×H×W×C ( C = 3 ) . Following ViT [ 13 ] , each frame is divided into K ×K non-overlapping patches which are then mapped into visual tokens using a linear embedding layer E ∈ R3K2×d . Since self-attention is permutation invariant , in order to preserve the information regarding the location of each patch within space and time we also learn two positional embeddings , one for space : ps ∈ R1×S×d and one for time : pt ∈ RT×1×d . These are then added to the initial visual tokens . Finally , the token sequence is processed by L Transformer layers . The visual token at layer l , spatial location s and temporal location t is denoted as : zls , t ∈ Rd , l = 0 , . . . , L− 1 , s = 0 , . . . , S − 1 , t = 0 , . . . , T − 1 . ( 1 ) In addition to the ST visual tokens extracted from the video , a special classification token zlcls ∈ Rd is prepended to the token sequence [ 12 ] . The l−th Transformer layer processes the visual tokens Zl ∈ R ( ST+1 ) ×d of the previous layer using a series of Multi-head Self-Attention ( MSA ) , Layer Normalization ( LN ) , and MLP ( Rd → R4d → Rd ) layers as follows : Yl = MSA ( LN ( Zl−1 ) ) + Zl−1 , ( 2 ) Zl = MLP ( LN ( Yl ) ) +Yl . ( 3 ) The main computation of a single full space-time Self-Attention ( SA ) head boils down to calculating : yls , t = T−1∑ t′=0 S−1∑ s′=0 Softmax { ( qls , t · kls′ , t′ ) / √ dh } vls′ , t′ , { s=0 , ... , S−1 t=0 , ... , T−1 } ( 4 ) where qls , t , k l s , t , v l s , t ∈ Rdh are the query , key , and value vectors computed from zls , t ( after LN ) using embedding matrices Wq , Wk , Wv ∈ Rd×dh . Finally , the output of the h heads is concatenated and projected using embedding matrix Wh ∈ Rhdh×d . The complexity of the full model is : O ( 3hTSddh ) ( qkv projections ) +O ( 2hT 2S2dh ) ( MSA for h attention heads ) +O ( TS ( hdh ) d ) ( multi-head projection ) +O ( 4TSd2 ) ( MLP ) 2 . From these terms , our goal is to reduce the costO ( 2T 2S2dh ) ( for a single attention head ) of the full space-time attention which is the dominant term 3 . For clarity , from now on , we will drop constant terms and dh to report complexity unless necessary . Hence , the complexity of the full space-time attention is O ( T 2S2 ) . Our baseline is a model that performs a simple approximation to the full space-time attention by applying , at each Transformer layer , spatial-only attention : yls , t = S−1∑ s′=0 Softmax { ( qls , t · kls′ , t ) / √ dh } vls′ , t , { s=0 , ... , S−1 t=0 , ... , T−1 } ( 5 ) 1Other attempts of exploiting local attention can be found in [ 29 , 7 , 49 ] , however they are also different in scope , task/domain and implementation . 2For this work , we used S = 196 , T = { 8 , 16 , 32 } and d = 768 ( for a ViT-B backbone ) . 3The MLP complexity is by no means negligible , however the focus of this work ( similarly to [ 3 , 1 ] ) is on reducing the complexity of the self-attention component . the complexity of which is O ( TS2 ) . Notably , the complexity of the proposed space-time mixing attention is also O ( TS2 ) . Following spatial-only attention , simple temporal averaging is performed on the class tokens zfinal = 1T ∑ t zL−1t , cls to obtain a single feature that is fed to the linear classifier . Recent work by [ 3 , 1 ] has focused on reducing the cost O ( T 2S2 ) of the full space-time attention of Eq . 4 . Bertasius et al . [ 3 ] proposed the factorised attention : ỹls , t = T−1∑ t′=0 Softmax { ( qls , t · kls , t′ ) / √ dh } vls , t′ , yls , t = S−1∑ s′=0 Softmax { q̃ls , t · k̃ls′ , t ) / √ dh } ṽls′ , t , { s = 0 , . . . , S − 1 t = 0 , . . . , T − 1 } , ( 6 ) where q̃ls , t , k̃ l s′ , tṽ l s′ , t are new query , key and value vectors calculated from ỹ l s , t 4 . The above model reduces complexity to O ( T 2S + TS2 ) . However , temporal attention is performed for a fixed spatial location which is ineffective when there is camera or object motion and there is spatial misalignment between frames . The work of [ 1 ] is concurrent to ours and proposes the following approximation : Ls Transformer layers perform spatial-only attention as in Eq . 5 ( each with complexity O ( S2 ) ) . Following this , there are Lt Transformer layers performing temporal-only attention on the class tokens zLst . The complexity of the temporal-only attention is , in general , O ( T 2 ) . Our model aims to better approximate the full space-time self-attention ( SA ) of Eq . 4 while keeping complexity to O ( TS2 ) , i.e . inducing no further complexity to a spatial-only model . To achieve this , we make a first approximation to perform full space-time attention but restricted to a local temporal window [ −tw , tw ] : yls , t = t+tw∑ t′=t−tw S−1∑ s′=0 Softmax { ( qls , t · kls′ , t′ ) / √ dh } vls′ , t′ = t+tw∑ t′=t−tw Vlt′a l t′ , { s=0 , ... , S−1 t=0 , ... , T−1 } ( 7 ) where Vlt′ = [ v l 0 , t′ ; v l 1 , t′ ; . . . ; v l S−1 , t′ ] ∈ Rdh×S and alt′ = [ al0 , t′ , al1 , t′ , . . . , alS−1 , t′ ] ∈ RS is the vector with the corresponding attention weights . Eq . 7 shows that , for a single Transformer layer , yls , t is a spatio-temporal combination of the visual tokens in the local window [ −tw , tw ] . It follows that , after k Transformer layers , yl+ks , t will be a spatio-temporal combination of the visual tokens in the local window [ −ktw , ktw ] which in turn conveniently allows to perform spatio-temporal attention over the whole clip . For example , for tw = 1 and k = 4 , the local window becomes [ −4 , 4 ] which spans the whole video clip for the typical case T = 8 . The complexity of the local self-attention of Eq . 7 is O ( T ( 2tw + 1 ) 2S2 ) . To reduce this even further , we make a second approximation on top of the first one as follows : the attention between spatial locations s and s′ according to the model of Eq . 7 is : t+tw∑ t′=t−tw Softmax { ( qls , t · kls′ , t′ ) / √ dh } vls′ , t′ , ( 8 ) i.e . it requires the calculation of 2tw +1 attentions , one per temporal location over [ −tw , tw ] . Instead , we propose to calculate a single attention over [ −tw , tw ] which can be achieved by qls , t attending kls′ , −tw : tw , [ k l s′ , t−tw ; . . . ; k l s′ , t+tw ] ∈ R ( 2tw+1 ) dh . Note that to match the dimensions of qls , t and kls′ , −tw : tw a further projection of k l s′ , −tw : tw to R dh is normally required which has complexity O ( ( 2tw +1 ) d 2 h ) and hence compromises the goal of an efficient implementation . To alleviate this we use the “ shift trick ” [ 45 , 26 ] which allows to perform both zero-cost dimensionality reduction , spacetime mixing and attention ( between qls , t and k l s′ , −tw : tw ) in O ( dh ) . In particular , each t ′ ∈ [ −tw , tw ] is assigned dt ′ h channels from dh ( i.e . ∑ t′ d t′ h = dh ) . Let k l s′ , t′ ( d t′ h ) ∈ Rd t′ h denote the operator for 4More precisely , Eq . 6 holds for h = 1 heads . For h > 1 , the different heads ỹl , hs , t are concatenated and projected to produce ỹls , t . indexing the dt ′ h channels from k l s′ , t′ . Then , a new key vector is constructed as : k̃ls′ , −tw : tw , [ k l s′ , t−tw ( d t−tw h ) , . . . , k l s′ , t+tw ( d t+tw h ) ] ∈ R dh . ( 9 ) Fig . 2 shows how the key vector k̃ls′ , −tw : tw is constructed . In a similar way , we also construct a new value vector ṽls′ , −tw : tw . Finally , the proposed approximation to the full space-time attention is given by : ylss , t = S−1∑ s′=0 Softmax { ( qlss , t · k̃ls′ , −tw : tw/ √ dh } ṽls′ , −tw : tw , { s=0 , ... , S−1 t=0 , ... , T−1 } . ( 10 ) This has the complexity of a spatial-only attention ( O ( TS2 ) ) and hence it is more efficient than previously proposed video transformers [ 3 , 1 ] . Our model also provides a better approximation to the full space-time attention and as shown by our results it significantly outperforms [ 3 , 1 ] . Temporal Attention aggregation : The final set of the class tokens zL−1t , cls , 0 ≤ t ≤ L− 1 are used to generate the predictions . To this end , we propose to consider the following options : ( a ) simple temporal averaging zfinal = 1T ∑ t z L−1 t , cls as in the case of our baseline . ( b ) An obvious limitation of temporal averaging is that the output is treated purely as an ensemble of per-frame features and , hence , completely ignores the temporal ordering between them . To address this , we propose to use a lightweight Temporal Attention ( TA ) mechanism that will attend to the T classification tokens . In particular a zfinal token attends the sequence [ zL−10 , cls , . . . , z L−1 T−1 , cls ] using a temporal Transformer layer and then fed as input to the classifier . This is akin to the ( concurrent ) work of [ 1 ] with the difference being that in our model we found that a single TA layer suffices whereas [ 1 ] uses Lt. A consequence of this is that the complexity of our layer is O ( T ) vs O ( 2 ( Lt − 1 ) T 2 + T ) of [ 1 ] . Summary token : As an alternative to TA , herein , we also propose a simple lightweight mechanism for information exchange between different frames at intermediate layers of the network . Given the set of tokens for each frame t , Zl−1t ∈ R ( S+1 ) ×dh ( constructed by concatenating all tokens zl−1s , t , s = 0 , . . . , S ) , we compute a new set of R tokens Z l r , t = φ ( Z l−1 t ) ∈ RR×dh which summarize the frame information and hence are named “ Summary ” tokens . These are then , appended to the visual tokens of all frames to calculate the keys and values so that the query vectors attend the original keys plus the Summary tokens . Herein , we explore the case that φ ( . ) performs simple spatial averaging zl0 , t = 1 S ∑ s z l s , t over the tokens of each frame ( R = 1 for this case ) . Note that , forR = 1 , the extra cost that the Summary token induces is O ( TS ) . X-ViT : We call the Video Transformer based on the proposed ( a ) space-time mixing attention and ( b ) lightweight global temporal attention ( or summary token ) as X-ViT .

Summary: In this paper, the authors focus on reducing the computation cost caused by the space-time attention modeling in the Transformer framework for video recognition. Specifically, the proposal includes 1) a space-time mixing attention module that mixes channels from neighboring frames within a local window as key-value tokens, and 2) a temporal attention layer at the end of the model to aggregate global temporal information. The authors provide promising results on Kinetics-400 and SSv2 in the paper.