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: This paper proposes a new computation block that unifies 3D convolution and transformer block for video action recognition. Starting from the transformer block, the core of the proposed block is its attention module which is equivalent to 3D CNN or transformer block depending on the design of the learnable parameters. On the standard action recognition benchmark, the proposed method outperforms prior art in both accuracy and computation efficiency. The comprehensive ablation studies provide insights on the design choices.