UniFormer: Unified Transformer for Efficient Spatial-Temporal Representation Learning

1 INTRODUCTION . Learning spatiotemporal representations is a fundamental task for video understanding . Basically , there are two distinct challenges . On the one hand , videos contain large spatiotemporal redundancy , where target motions across local neighboring frames are subtle . On the other hand , videos contain complex spatiotemporal dependency , since target relations across long-range frames are dynamic . The advances in video classification have mostly driven by 3D convolutional neural networks ( Tran et al. , 2015 ; Carreira & Zisserman , 2017b ; Feichtenhofer et al. , 2019 ) and spatiotemporal transformers ( Bertasius et al. , 2021 ; Arnab et al. , 2021 ) . Unfortunately , each of these two frameworks focuses on one of the aforementioned challenges . 3D convolution can capture detailed and local spatiotemporal features , by processing each pixel with context from a small 3D neighborhood ( e.g. , 3×3×3 ) . * Equally-contributed first authors ( { kc.li , yl.wang } @ siat.ac.cn ) †Corresponding author ( qiaoyu @ pjlab.org.cn ) Hence , it can reduce spatiotemporal redundancy across adjacent frames . However , due to the limited receptive field , 3D convolution suffers from difficulty in learning long-range dependency ( Wang et al. , 2018 ; Li et al. , 2020a ) . Alternatively , vision transformers are good at capturing global dependency , with the help of self-attention among visual tokens ( Dosovitskiy et al. , 2021 ) . Recently , this design has been introduced in video classification via spatiotemporal attention mechanism ( Bertasius et al. , 2021 ) . However , we observe that , video transformers are often inefficient to encode local spatiotemporal features in the shallow layers . We take the well-known and typical TimeSformer ( Bertasius et al. , 2021 ) for illustration . As shown in Figure 1 , TimeSformer indeed learns detailed video representations in the early layers , but with very redundant spatial and temporal attention . Specifically , spatial attention mainly focuses on the neighbor tokens ( mostly in 3×3 local regions ) , while learning nothing from the rest tokens in the same frame . Similarly , temporal attention mostly only aggregates tokens in the adjacent frames , while ignoring the rest in the distant frames . More importantly , such local representations are learned from global token-to-token similarity comparison in all layers , requiring large computation cost . This fact clearly deteriorates computation-accuracy balance of such video transformer ( Figure 2 ) . To tackle these difficulties , we propose to effectively unify 3D convolution and spatiotemporal selfattention in a concise transformer format , thus we name the network Unified transFormer ( UniFormer ) , which can achieve a preferable balance between efficiency and effectiveness . More specifically , our UniFormer consists of three core modules , i.e. , Dynamic Position Embedding ( DPE ) , Multi-Head Relation Aggregator ( MHRA ) , and Feed-Forward Network ( FFN ) . The key difference between our UniFormer and traditional video transformers is the distinct design of our relation aggregator . First , instead of utilizing a self-attention mechanism in all layers , our proposed relation aggregator tackles video redundancy and dependency respectively . In the shallow layers , our aggregator learns local relation with a small learnable parameter matrix , which can largely reduce computation burden by aggregating context from adjacent tokens in a small 3D neighborhood . In the deep layers , our aggregator learns global relation with similarity comparison , which can flexibly build long-range token dependencies from distant frames in the video . Second , different from spatial and temporal attention separation in the traditional transformers ( Bertasius et al. , 2021 ; Arnab et al. , 2021 ) , our relation aggregator jointly encodes spatiotemporal context in all the layers , which can further boost video representations in a joint learning manner . Finally , we build up our model by progressively integrating UniFormer blocks in a hierarchical manner . In this case , we enlarge the cooperative power of local and global UniFormer blocks for efficient spatiotemporal representation learning in videos . We conduct extensive experiments on the popular video benchmarks , e.g. , Kineticss-400 ( Carreira & Zisserman , 2017a ) , Kinetics-600 ( Carreira et al. , 2018 ) and SomethingSomething V1 & V2 ( Goyal et al. , 2017b ) . With only ImageNet-1K pretraining , our UniFormer achieves 82.9 % /84.8 % top-1 accuracy on Kinetics-400/Kinetics-600 , while requiring 10× fewer GFLOPs than other comparable methods ( e.g. , 16.7× fewer GFLOPs than ViViT ( Arnab et al. , 2021 ) with JFT-300M pre-training ) . For Something-Something V1 and V2 , our UniFormer achieves 60.9 % and 71.2 % top-1 accuracy respectively , which are new state-of-the-art performances . 2 RELATED WORK . Convolution-based Video Networks . 3D Convolution Neural Networks ( CNNs ) have been dominant in video understanding ( Tran et al. , 2015 ; Feichtenhofer et al. , 2019 ) . However , they suffer from the difficult optimization problem and large computation cost . To resolve this issue , I3D ( Carreira & Zisserman , 2017b ) inflates the pre-trained 2D convolution kernels for better optimization . Other prior works ( Tran et al. , 2018 ; Qiu et al. , 2017 ; Tran et al. , 2019 ; Feichtenhofer , 2020 ; Wang et al. , 2020a ) try to factorize 3D convolution kernel in different dimensions to reduce complexity . Recent methods propose well-designed modules to enhance the temporal modeling ability for 2D CNNs ( Wang et al. , 2016 ; Lin et al. , 2019 ; Luo & Yuille , 2019 ; Jiang et al. , 2019 ; Liu et al. , 2020a ; Li et al. , 2020b ; Kwon et al. , 2020 ; Li et al. , 2020a ; 2021a ; Wang et al. , 2020b ) . However , 3D convolution struggles to capture long-range dependency , due to the limited receptive field . Transformer-based Video Networks . Vision Transformers ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021a ; b ; Liu et al. , 2021a ) have been popular for vision tasks and outperform many CNNs . Based on ViT , several prior works ( Bertasius et al. , 2021 ; Neimark et al. , 2021 ; Sharir et al. , 2021 ; Li et al. , 2021b ; Arnab et al. , 2021 ; Bulat et al. , 2021 ; Patrick et al. , 2021 ; Zha et al. , 2021 ) propose different variants for spatiotemporal learning , verifying the outstanding ability of the transformer to capture long-term dependencies . To reduce high dot-product computation , MViT ( Fan et al. , 2021 ) introduces the hierarchical structure and pooling self-attention , while Video Swin ( Liu et al. , 2021b ) advocates an inductive bias of locality for video . Nevertheless , the self-attention mechanism is inefficient to encode low-level features , hindering their high potential . To tackle this challenge , different from Video Swin that applies self-attention in a local 3D window , we adopt 3D convolution in a concise transformer format to encode local features . Besides , we follow their hierarchical designs and propose our UniFormer , achieving powerful performance for video understanding . 3 METHOD . In this section , we describe our UniFormer in detail . First , we introduce the overall architecture of the UniFormer block . Then , we explain the vital designs of our UniFormer for spatiotemporal modeling , i.e. , multi-head relation aggregator and dynamic position embedding . Finally , we hierarchically stack UniFormer blocks to build up our video network . 3.1 OVERVIEW OF UNIFORMER BLOCK . To overcome problems of spatiotemporal redundancy and dependency , we propose a novel and concise Unified transFormer ( UniFormer ) shown in Figure 3 . We utilize a basic transformer format ( Vaswani et al. , 2017 ) but specially design it for efficient and effective spatiotemporal representation learning . Specifically , our UniFormer block consists of three key modules : Dynamic Position Embedding ( DPE ) , Multi-Head Relation Aggregator ( MHRA ) , and Feed-Forward Network ( FFN ) : X = DPE ( Xin ) +Xin , ( 1 ) Y = MHRA ( Norm ( X ) ) +X , ( 2 ) Z = FFN ( Norm ( Y ) ) +Y . ( 3 ) Considering the input token tensor ( frame volumes ) Xin ∈ RC×T×H×W , we first introduce DPE to dynamically integrate 3D position information into all the tokens ( Eq . 1 ) , which effectively makes use of spatiotemporal order of the tokens for video modeling . Then , we leverage MHRA to aggregate each token with its contextual tokens ( Eq . 2 ) . Different from the regular Multi-Head Self-Attention ( MHSA ) , our MHRA smartly tackles local video redundancy and global video dependency , by flexible designs of token affinity learning in the shallow and deep layers . Finally , we add FFN with two linear layers for pointwise enhancement of each token ( Eq . 3 ) . 3.2 MULTI-HEAD RELATION AGGREGATOR . As discussed above , we should solve large local redundancy and complex global dependency , for efficient and effective spatiotemporal representation learning . Unfortunately , the popular 3D CNNs and spatiotemporal transformers only focus on one of these two challenges . For this reason , we design an alternative Relation Aggregator ( RA ) , which can flexibly unify 3D convolution and spatiotemporal self-attention in a concise transformer format , solving video redundancy and dependency in the shallow layers and deep layers respectively . Specifically , our MHRA conducts token relation learning via multi-head fusion : Rn ( X ) = AnVn ( X ) , ( 4 ) MHRA ( X ) = Concat ( R1 ( X ) ; R2 ( X ) ; · · · ; RN ( X ) ) U . ( 5 ) Given the input tensor X ∈ RC×T×H×W , we first reshape it to a sequence of tokens X ∈ RL×C , L=T×H×W . Rn ( · ) is the relation aggregator ( RA ) in the n-th head , and U ∈ RC×C is a learnable parameter matrix to integrate N heads . Moreover , each RA consists of token context encoding and token affinity learning . Via a linear transformation , one can transform the original token into context Vn ( X ) ∈ RL× C N . Subsequently , the relation aggregator can summarize context with the guidance of the token affinity An ∈ RL×L . The key in our RA is how to learn An in videos . Local MHRA . In the shallow layers , we aim at learning detailed video representation from the local spatiotemporal context in small 3D neighborhoods . This coincidentally shares a similar insight with the design of a 3D convolution filter . As a result , we design the token affinity as a learnable parameter matrix operated in the local 3D neighborhood , i.e. , given one anchor token Xi , RA learns local spatiotemporal affinity between this token and other tokens in the small tube Ωt×h×wi : Alocaln ( Xi , Xj ) = a i−j n , where j ∈ Ωt×h×wi , ( 6 ) where an ∈ Rt×h×w and Xj refers to any neighbor token in Ωt×h×wi . ( i − j ) means the relative token index that determines the aggregating weight ( Appendix A shows more details ) . In the shallow layers , video contents between adjacent tokens vary subtly , it is significant to encode detailed features with local operator to reduce redundancy . Hence , the token affinity is designed as a local learnable parameter matrix , whose values only depend on the relative 3D position between tokens . Comparison to 3D Convolution Block . Interestingly , we find that our local MHRA can be interpreted as a spatiotemporal extension of MobileNet block ( Sandler et al. , 2018 ; Tran et al. , 2019 ; Feichtenhofer , 2020 ) . Specifically , the linear transformation V ( · ) can be instantiated as pointwise convolution ( PWConv ) . Furthermore , the local token affinity Alocaln is a spatiotemporal matrix that operated on each output channel ( or head ) Vn ( X ) , thus the relation aggregator Rn ( X ) = Alocaln Vn ( X ) can be explained as a depthwise convolution ( DWConv ) . Finally , all heads are concatenated and fused by a linear matrix U , which can also be instantiated as pointwise convolution ( PWConv ) . As a result , this local MHRA can be reformulated with a manner of PWConv-DWConv-PWConv in the MobileNet block . In our experiments , we flexibly instantiate our local MHRA as such channelseparated spatiotemporal convolution , so that our UniFormer can inherit computation efficiency for light-weight video classification . Different from the MobileNet block , our UniFormer block is designed as a generic transformer format , thus an extra FFN is inserted after MHRA , which can further mix token context at each spatiotemporal position to boost classification accuracy . Global MHRA . In the deep layers , we focus on capturing long-term token dependency in the global video clip . This naturally shares a similar insight with the design of self-attention . Hence , we design the token affinity via comparing content similarity among all the tokens in the global view : Aglobaln ( Xi , Xj ) = eQn ( Xi ) TKn ( Xj ) ∑ j′∈ΩT×H×W e Qn ( Xi ) TKn ( Xj′ ) , ( 7 ) where Xj can be any token in the global 3D tube with size of T×H×W , while Qn ( · ) and Kn ( · ) are two different linear transformations . Most video transformers apply self-attention in all stages , introducing a large amount of calculation . To reduce the dot-product computation , the prior works tend to divide spatial and temporal attention ( Bertasius et al. , 2021 ; Arnab et al. , 2021 ) , but it deteriorates the spatiotemporal relation among tokens . In contrast , our MHRA performs local relation aggregation in the early layers , which largely saves the computation of token comparison . Hence , instead of factorizing spatiotemporal attention , we jointly encode spatiotemporal relation in our MHRA for all the stages , in order to achieve a preferable computation-accuracy balance . Comparison to Transformer Block . In the deep layers , our UniFormer block is equipped with a global MHRA Aglobaln ( Eq . 7 ) . It can be instantiated as a spatiotemporal self attention , where Qn ( · ) , Kn ( · ) and Vn ( · ) become Query , Key and Value in the transformer ( Dosovitskiy et al. , 2021 ) . Hence , it can effectively learn long-term dependency . Instead of spatial and temporal factorization in the previous video transformers ( Bertasius et al. , 2021 ; Arnab et al. , 2021 ) , our global MHRA is based on joint spatiotemporal learning to generate more discriminative video representation . Moreover , we adopt dynamic position embedding ( DPE , see Section 3.3 ) to overcome permutation-invariance , which can maintain translation-invariance and is friendly to different input clip lengths .

Summary: The paper proposes a UniFormer architecture, which combines 3D convolutions and spatiotemporal self-attention for efficient and effective video classification performance. Specifically, the authors show that using local 3D convolution-like layers in the early stages of the network and global self-attention layers in the later stages of the network provides a good tradeoff between efficiency and accuracy. The proposed method reports state-of-the-art results on several major action recognition benchmarks.