Shifted Chunk Transformer for Spatio-Temporal Representational Learning

1 Introduction . Spatio-temporal representational learning tries to model complicated intra-frame and inter-frame relationships , and it is critical to various tasks such as action recognition [ 20 ] , action detection [ 55 , 52 ] , object tracking [ 23 ] , and action anticipation [ 21 ] . Deep learning based spatio-temporal representation learning approaches have been largely explored since the success of AlexNet on image classification [ 24 , 11 ] . Previous deep spatio-temporal learning can be mainly divided into two aspects : deep ConvNets based methods [ 35 , 15 , 16 ] and deep sequential learning based methods [ 55 , 28 , 29 ] . Deep ConvNets based methods are primarily integrated various factorization techniques [ 51 , 34 ] , or a priori [ 16 ] for efficient spatio-temporal learning [ 15 ] . Some works focus on extracting effective spatio-temporal features [ 41 , 8 ] or capturing complicated long-range dependencies [ 49 ] . Deep sequential learning based methods try to formulate the spatial and temporal relationships through advanced deep sequential models [ 28 ] or the attention mechanism [ 29 ] . On the other hand , the Transformer has become the de-facto standard for sequential learning tasks such as speech and language processing [ 44 , 12 , 54 , 19 ] . The great success of Transformer on natural language processing ( NLP ) has inspired computer vision community to explore self-attention 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . structures for several vision tasks , e.g. , image classification [ 13 , 40 , 36 ] , object detection [ 6 ] , and super-resolution [ 33 ] . The main difficulty in pure-Transformer models for vision is that Transformers lack the inductive biases of convolutions , such as translation equivariance , and they require more data [ 13 ] or stronger regularisation [ 40 ] in the training . It is only very recently that , vision Transform ( ViT ) , a pure Transformer architecture , has outperformed its convolutional counterparts in image classification when pre-trained on large amounts of data [ 13 ] . However , the hurdle is aggravated when the pure-Transformer design is applied to spatio-temporal representational learning . Recently , a few attempts have been made to design pure-Transformer structures for spatio-temporal representation learning [ 4 , 5 , 14 , 2 ] . Simply applying Transformer to 3D video domain is computationally intensive [ 4 ] . The Transformer based spatio-temporal learning methods primarily focus on designing efficient variants by factorization along spatial- and temporal-dimensions [ 4 , 5 ] , or employing a multi-scale pyramid structure for a trade-off between the resolution and channel capacity while reducing the memory and computational cost [ 14 ] . The spatio-temporal learning capacity can be further improved by extracting more effective fine-grained features through advanced and efficient intra-frame and inter-frame representational learning . In this work , we propose a novel spatio-temporal learning framework based on pure-Transformer , called shifted chunk Transformer as illustrated in Fig . 1 , which extracts effective fine-grained intraframe features with a low computational complexity leveraging the recent advance of Transformer in NLP [ 22 ] . Specially , we divide each frame into several local windows called image chunks , and construct a hierarchical image chunk Transformer , which employs locality-sensitive hashing ( LSH ) to enhance the dot-product attention in each chunk and reduces the memory and computation consumption significantly . To fully consider the motion effect of object , we design a robust selfattention module , shifted self-attention , which explicitly extracts correlations from nearby frames . We further design a pure-Transformer based frame-wise attention module , clip encoder , to model the complicated inter-frame relationships with a minimal extra computational cost . Our contributions can be summarized as follows : • We construct an image chunk self-attention to mine fine-grained intra-frame features leveraging the recent advance of Transformer . The hierarchical image chunk Transformer employs locality-sensitive hashing ( LSH ) [ 3 ] to reduce the memory and computation consumption significantly , which enables an effective spatio-temporal learning directly from a tiny patch . • We build a shifted self-attention to fully consider the motion effect of objects , which yields effective modeling of complicated inter-frame variances in the spatio-temporal representa- tional learning . Furthermore , a clip encoder with a pure-Transformer structure is employed for frame-wise attention , which models complicated and long-term inter-frame relationships at a minimal extra cost . • The shifted chunk Transformer with pure-Transformer outperforms previous state-of-theart approaches on several action recognition benchmarks , including Kinectics-400 [ 20 ] , Kinetics-600 [ 7 ] , UCF101 [ 39 ] and HMDB51 [ 25 ] . 2 Related Work . Conventional deep learning based action recognition Conventional deep spatio-temporal representational learning mainly involves two aspects : deep sequential learning based methods [ 55 , 28 , 29 ] and deep ConvNet based methods [ 35 , 15 , 16 ] . The recurrent networks can be extended to 3D spatio-temporal domain for action recognition [ 28 ] . In deep ConvNet based methods , two-stream ConvNet employs two branches of 2D ConvNets and explicitly models motion by optical flow [ 38 ] . The C3D [ 41 ] and I3D [ 8 ] directly extend 2D ConvNets to 3D ConvNets , which is natural for 3D spatio-temporal representational learning [ 9 ] . However , the 3D ConvNet requires significantly more computation and more training data to achieve a desired accuracy [ 51 ] . Thus , P3D [ 34 ] and S3D [ 51 ] attempt to factorize the 3D convolution into a 2D spatial convolution and a 1D temporal convolution . SlowFast network [ 16 ] and X3D [ 15 ] conduct trade-offs among resolution , temporal frame rate and the number of channels for the efficient video recognition . Non-local network [ 49 ] proposes to add non-local operations in deep network and captures long-range dependencies . The recent pure-Transformer based spatio-temporal learning enables longer dependency relationship modeling and further increases the accuracy of action recognition [ 4 ] . Vision Transformers NLP community has witnessed the great success of pre-training by Transformer [ 44 , 12 ] , and it has been emerging for image classification [ 13 , 40 , 36 ] , object detection [ 6 ] , and image super-resolution [ 33 ] . CPVT [ 10 ] employs pre-defined and independent input tokens to increase the generalization for image classification . Pure-Transformer network has no inductive bias or prior as ConvNets . ViT [ 13 ] pre-trains on large amounts of data and attains excellent results on image classification . CvT [ 50 ] introduces convolutions into ViT to yield better performance and efficiency . DeiT [ 40 ] and MViT [ 14 ] instead employ distillation and multi-scale to cope with the training difficulty . PVT [ 47 ] and segmentation Transformer [ 53 ] further extend Transformer to dense prediction tasks , e.g. , object detection and semantic segmentation . Simply applying Transformer to 3D video spatio-temporal representational learning aggravates the training difficulty significantly , and it requires advanced model design for pure-Transformer based spatio-temporal learning . Transformer based action recognition Recently , only a few works have been conducted using pure-Transformer for spatio-temporal learning [ 4 , 14 , 5 , 2 ] . Most of the efforts focus on designing efficient Transformer models to reduce the computation and memory consumption . ViViT [ 4 ] and TimeSformer [ 5 ] study various factorization methods along spatial- and temporal-dimensions . MViT [ 14 ] conducts a trade-off between resolution and the number of channels , and constructs a multi-scale Transformer to learn a hierarchy from simple dense resolution and fine-grained features to complex coarse features . VATT [ 2 ] conducts unsupervised multi-modality self-supervised learning with a pure-Transformer structure . In this work , we extract fine-grained intra-frame features from each tiny patch and model complicated inter-frame relationship through efficient and advanced self-attention blocks . 3 Method . In this section , we describe each component of our shifted chunk Transformer for spatio-temporal representation learning in video based action recognition . 3.1 Overview . Let X = [ X1 , X2 , · · · , XT ] ∈ RT×H×W×3 be one input clip of T RGB frames sampled from a video , where Xi ∈ RH×W×3 is the i-th frame in the clip , and H and W are the frame size . To design an efficient pure-Transformer based spatio-temporal learning , we construct a shifted chunk Transformer , including image chunk self-attention blocks , shifted multi-head self-attention blocks , and a clip encoder , as illustrated in Fig . 1 . We first construct an image chunk self-attention block leveraging the advanced efficient Transformer design in NLP [ 22 ] , which is illustrated in the left of Fig . 1 . The locality-sensitive hashing ( LSH ) [ 3 ] in the image chunk self-attention enables a relatively small patch as a token , thus it is capable of extracting fine-grained intra-frame features . A linear pooling layer is designed to adaptively reduce the resolution after LSH attention . After that , a shifted self-attention is designed to extract motion related inter-frame features . Our shifted self-attention considers the motion of objects across nearby frames and explicitly models the temporal relationship into self-attention . The frame encoder shown in dark grey color can be an effective feature extractor which can be stacked for several times . The hierarchical frame encoder and image chunk self-attention further learns an effective multi-level feature from local to global abstraction . Lastly , we employ a pure-Transformer to learn complicated inter-frame relationships and frame-wise attention along the temporal dimension . We use multi-head self-attention ( MSA ) in all the blocks . 3.2 Image Chunk Self-Attention Transformer can learn complicated long range dependencies which can be computed through the high efficient matrix product [ 44 ] . Different from convolution which has inherent inductive bias [ 13 ] , Transformer learns the entire features from data . The main challenge for a pureTransformer based vision model mainly involves two aspects : 1 ) how to design an efficient model to learn effective features from the entire image , because simply treating each pixel as a token is computationally intensive , 2 ) how to train this powerful model and learn various effective features from data . ViT [ 13 ] treats each patch of size 16× 16 as one token and pre-trains the model with large amounts of data . We argue that Transformer with a smaller patch as a token can extract fine-grained features which improves spatio-temporal learning for action recognition . For ViT-B-16 [ 13 ] of crop size 224× 224 as a frame encoder followed by a shifted MSA and a clip encoder , a smaller patch size of 8× 8 yields better accuracy on Kinectics-400 and UCF101 as shown in Table 1 . The Transformer computes each pairwise correlation through a dot product , thus it has a high computational complexity of O ( L2 ) , where L is the totally number of tokens . In natural language processing ( NLP ) , LSH attention [ 22 ] employs locality-sensitive hashing ( LSH ) bucketing [ 3 ] and chunk sorting for queries and keys to approximate the attention matrix computation . Leveraging the efficient LSH approximation , the LSH attention reduce the computation complexity to O ( L logL ) . To preserve locality property and learn transition and rotation invariant low-level features from images , we firstly design a visual local transformer ( ViLT ) with shared parameters along different image local windows , or image chunks . Each image chunk consists of multiple tiny patches as illustrated in the left bottom block of Fig . 1 . We intend to employ patches of a small size in the ViLT which is the first level abstraction of the input , so that the model extracts a fine-grained representation which enhances the entire spatial-temporal learning . Employing the small patch yields large number of tokens in the following self-attention , and we construct an image locality-sensitive hashing ( LSH ) attention leveraging advanced design of Transformer in NLP [ 22 ] . The image LSH attention can efficiently extract higher-level and plentiful intra-frame features ranging from a tiny patch to the entire image . The framework of image chunk self-attention is illustrated in the left part of Fig . 1 . Visual local transformer ( ViLT ) In the shifted chunk Transformer , we firstly construct a ViLT which slides one self-attention for each tiny patch along the whole image . The ViLT is illustrated in the bottom block of Fig . 1 . Let h , w be the height and width of the tiny image patch p ∈ Rh×w×3 . Following the success of ViT [ 13 ] , we treat each patch as one dimensional vector of length h × w × 3 . Suppose each chunk consists of m × n tiny patches , denoted as { p1,1 , p1,2 , · · · , p1 , n ; p2,1 , · · · , p2 , n ; · · · ; pm,1 , · · · , pm , n } . After flattening the chunk into a list of tiny patches , we denote the chunk as { p1 , p2 , · · · , pL } without loss of generality , where L = m× n. In ViLT , we use a learnable 1D position embedding Epos ∈ RL×D to retain position information z0 = [ p1E ; p2E ; · · · ; pm×nE ] +Epos , E ∈ R ( h×w×3 ) ×D , ( 1 ) where E is the linear patch embedding matrix , and D is the embedding dimension . Then we can construct alternating layers of multi-head self-attention ( MSA ) , MLP with GELU [ 18 ] non-linearity , Layernorm ( LN ) and residual connections [ 44 ] for the chunk as z′m = MSA ( LN ( zm−1 ) ) + zm−1 , zm = MLP ( LN ( z ′ m ) ) + z ′ m , m = 1 , · · · , M , ( 2 ) where M is the number of blocks . We conduct the ViLT sliding the entire image without overlapping . The parameters of the ViLT are shared among all the image chunks , which forces the chunk self-attention to learn translation and rotation invariant , and fine-grained features . The tiny patch-wise feature extraction preserves the locality property , which is a strong prior for natural images . After the ViLT , we obtain the extracted features for the entire image denoted as y = [ y1 ; y2 ; · · · ; yL ; · · · ; yL×L′ ] , where the entire image can be split into L′ = m′ × n′ = d Hh×me × d W w×ne chunks , and we conduct zero padding for the last chunks in each row and column . The ViLT forces to learn image locality features which is a desired property for a low-level feature extractor [ 26 ] . For pure-Transformer based vision system , it also reduces the memory consumption significantly because it restricts the correlation of one tiny patch within the local chunk . Therefore , the memory and computational complexity of dot-product attention in ViLT can be reduced to O ( L2 ) compared to the complexity of conventional self-attention O ( ( L′ × L ) 2 ) . Image locality-sensitive hashing ( LSH ) attention After ViLT blocks , we obtain local fine-grained features of length L′ ×L . Since the patch size is tiny , the total number of patches can be large , which leads to more difficult training than other vision Transformers [ 4 , 13 ] . On the other hand , the problem of finding nearest neighbors quickly in high-dimensional spaces can be solved by locality-sensitive hashing ( LSH ) , which hashes similar input items into the same “ buckets ” with high probability . In NLP , LSH attention [ 22 ] is proposed to handle quite long sequence data , which employs localitysensitive hashing ( LSH ) bucketing approximation and bucket sorting to reduce the computational complexity of matrix product between query and key in self-attention . In dot-product attention , the softmax activation function pushes the attention weights close to 1 or 0 , which means the attention matrix is typically sparse . The query and key can be approximated by locality-sensitive hashing ( LSH ) [ 3 ] to reduce the computational complexity . Furthermore , through bucketing sort , the attention matrix product can be accelerated by a chunk triangular matrix product , which has been validated by LSH attention [ 22 ] . The used multi-head image LSH attention can be constructed as y′ = MSA ( LSHAtt ( LN ( y ) ) ) + y , s = MLP ( LN ( y′ ) ) + y′ , ( 3 ) where LSH attention LSHAtt ( · ) employs angular distance to conduct LSH hashing [ 3 ] . The image LSH attention reduces the memory and time complexity to O ( L′ × L log ( L′ × L ) ) , compared with that of conventional dot-product attention O ( ( L′ × L ) 2 ) . The image LSH attention reduces the complexity significantly because the patch size is tiny and the number of tiny patches L′ × L are large . The image level LSH attention in the second level learns relatively global features from the first level ’ s local fine-grained features . The hierarchical feature learning from local to global has been validated as an effective principle for vision system design [ 26 , 45 ] . Inspired by hierarchical abstraction in ConvNets [ 26 ] , we construct a linear pooling layer which firstly conducts squeeze then employs linear projection for feature dimension reduction . The linear pooling adaptively squeezes the sequence length by 14 . Reshape : s = [ s1,1 , · · · , s1 , n×n′ ; · · · ; sm×m′,1 , · · · , sm×m′ , n×n′ ] , Squeeze : s′ = [ s′1,1 , · · · , s′1 , n×n′/2 ; · · · ; s ′ m×m′/2,1 , · · · , s ′ m×m′/2 , n×n′/2 ] , Linear : t = [ s′1,1Et , · · · , s′1 , n×n′/2Et ; · · · ; s ′ m×m′/2,1Et , · · · , s ′ m×m′/2 , n×n′/2Et ] , ( 4 ) where s′i , j in s ′ concatenates s2i−1,2j−1 , s2i−1,2j , s2i,2j−1 and s2i,2j from s , Et is the linear projection matrix to adaptively reduce the number of dimensions after squeeze by half . The reshape is to retain the spatial relationship of each patch , and the squeeze reduces the number of patches by 14 and enlarges the feature dimensions by four times . The linear pooling layer forces the model to learn high-level global features in the following layers .

Summary: The paper proposes a novel transformer-based architecture for video processing. The input tokens are small patches extracted for each frame. The model consists of 4 modules: First, the tokens are processed locally, using a transformer with restricted, local connectivity (each image is split into multiple chunks and only the patches from the same chunk exchange messages using the self-attention layer). Second, a global processing is conducted using an efficient transformer (that use LSH attention approximation) to capture long-range interactions. In order to take into account the motion effect, a shifted self-attention is introduced, modifying the classical self-attention layer to use keys computed from the previous time step. Similar to language modelling, a special token ([CLS]) is used to extract frame-level representations that are fed into a final transformer block to compute the final video representation.