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: This paper proposes a new Transformer architecture for video modeling. There are following several components with their respective design motivation: - First, each frame is split into a large number of tiny patches (e.g. 4x4 pixels) which are then linearly projected into token vectors. Then, the authors propose Image Local Chunk Attention to enforce locality in cross attention to ensure attention is only computed among neighboring patches in local neighborhoods. The authors argue that it's necessary to have small patches for good accuracy. And the Image Local Chunk Attention is therefore there to reduce computation costs brought by small patches and also enforce locality.  - With the features computed from Image Local Chunk Attention, the authors then apply a locality sensitive hashing (LSH) attention to approximate the true cross attention to increase efficiency.  - Now tokens are mixed with attention for each frame, the authors then propose Shifted Multi-Head Self-Attention that computes cross-attention across frames to model object motion and spatial variances.  - Finally, a special CLS token is extracted from each frame and passed into another clip encoder to produce a clip-level summarization, which is used for final classification tasks. Overall, the authors proposed a set of methods to model both intra and inter frame relationships in videos using transformers and achieved SOTA results on several video classification tasks like Kinetics, UCF101, HMDB51 and Moment-in-Time. The authors also conducted thorough ablation to dissect the contribution of each components.