TokenLearner: Adaptive Space-Time Tokenization for Videos

1 Introduction . Videos provide an abundance of visual information . Video understanding particularly requires employing effective spatial-temporal processing of frames to capture long-range interactions [ 5 , 37 , 21 , 17 , 24 , 12 , 34 , 20 , 25 , 1 ] . An important aspect of this understanding is how to quickly learn which parts of the input video stream are important , both spatially and temporally , and to focus computational resources on them . But what basic processing mechanism are able to do so successfully ? Recent advancements in Transformers demonstrate improved accuracy on vision classification tasks . For example , departing from standard convolutional approaches , the Vision Transformer ( ViT ) [ 9 ] treats the image as a sequence of patches , utilizing the Transformer architecture [ 39 ] similar to text understanding . Standard approaches for video recognition take videos as stacked images ( i.e. , a spacetime volume ) and tend to extend 2D neural architectures to 3D ( e.g. , 3D-ResNets [ 17 , 5 , 38 , 11 ] ) . Motivated by ViT , recent approaches [ 2 , 3 ] also extend Transformers for videos by creating 3D ‘ tubelet ’ video tokens with regular 3D-grids , which often result in computationally heavy models . There are often too many tokens to process , especially for longer videos . The main question addressed in this work is how to adaptively learn the representation from visual inputs to most effectively capture the spatial information for image frames and spatio-temporal interactions for videos . Here are our main ideas : The first key observation is we are able to learn to represent visual data by learning to ‘ tokenize ’ the representations . This is in contrast to previous approaches which used densely sampled tokens e.g. , 16x16 or 32x32 over a series of attention layers [ 9 , 3 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Specifically , we can learn to compute important regions in the input image/video , making the tokens adapt to the input data . We compute multiple spatial weight maps per frame with a spatial attention mechanism , and use it for the tokenization . The goal of these maps is to learn which areas are of importance . Here , each spatial weight map is multiplied with the input to form a ‘ token ’ , to be processed by the subsequent learning modules . Furthermore , we find that very few tokens may be sufficient for a visual understanding task . More specifically , we show that one can significantly reduce the computational budget of video Transformers , by utilizing 8-16 tokens as an intermediate frame representation ( instead of keeping 200∼500 ) . Our TokenLearner is able to reduce the number of total FLOPS by half , while maintaining or even increasing the classification accuracy . The approach is simple , efficient , and , as shown by the results , outperforms methods including both convolutional methods and previous space-time Transformer ones from prior art . In video understanding tasks , we establish new state-of-the-art numbers on Kinetics-400 , Kinetics-600 , Charades , and AViD datasets by outperforming prior models . 2 TokenLearner Modules for Adaptive Tokenization . In visual Transformer architectures such as ViT [ 9 ] , an input image is first tokenized by splitting it into small ( e.g. , 16x16 ) spatial patches , which are used as input to the model . Similarly , in recent video Transformer architectures , such as ViViT [ 2 ] and TimeSformer [ 3 ] , the video is tokenized by cutting the video into 2d spatial or 3d spatio-temporal cubes on a regular grid . Instead of processing fixed , tokenized inputs , our attention module learns the tokens that are to be used for the recognition task . We gain several important properties by doing so : ( 1 ) We enable the adaptive tokenization so that the tokens can be dynamically selected conditioned on the input . ( 2 ) This also effectively reduces the total number of tokens for the transformer , which is particularly beneficial considering that there are many tokens in videos ( e.g. , 14 × 14 × 64 ) and the computation is quadratic to the number of tokens . ( 3 ) Finally , we provide an ability for each subsequent layer to learn to rely on different space-time tokenizations , potentially allowing different layers to capture different aspects of the video . These dynamically and adaptively generated tokens can be used in standard transformer architectures such as ViT for images and ViViT for videos . 2.1 TokenLearner . Let X be an input tensor with a space-time shape : X ∈ RT×H×W×C where H ×W corresponds to the spatial dimension of the input , T is the temporal dimension ( i.e. , number of frames ) , and C is the number of channels . Let Xt be a temporal slice of it , corresponding to the frame t : Xt ∈ RH×W×C . In the case of an image input , T = 1 and X = Xt . Note that X could also be an intermediate representation within a network , and Xt will be its slice in such case . For every time frame t , we learn to generate a series of S tokens , Zt = [ zi ] Si=1 , from the input frame Xt . Specifically , we formulate a tokenizer function , zi = Ai ( Xt ) , which maps the input frame Xt to a token vector zi : RH×W×C 7→ RC . The idea is to learn our tokenizer function Ai to adaptively select an informative combination of pixels ( or spatial locations ) in Xt , and we have S number of such functions . This way , our tokens will not be fixed splits of the input tensor , but a set of adaptively changing spatial selections . Different tokens will be mined per frame , allowing us to model their space-time relations/interactions in case of videos . We also set S to be smaller than H ×W ( e.g. , S = 8 and H ×W = 14× 14 ) , enabling the model to significantly reduce the computations needed for the layers following this module . Here , our tokenizer zi = Ai ( Xt ) is implemented with a spatial attention mechanism : i.e. , the model learns to compute a weight map ( of size H ×W ) conditioned on the input Xt , and is multiplied with Xt itself . More specifically , let αi ( Xt ) be a function generating the spatial H ×W × 1 weight map . Each token zi is generated by zi = Ai ( Xt ) = ρ ( Xt Aiw ) = ρ ( Xt γ ( αi ( Xt ) ) ) , ( 1 ) where is the Hadamard product ( i.e. , element-wise multiplication ) and Aiw ∈ RH×W×C is an intermediate weight tensor computed with the function αi ( Xt ) and the broadcasting function γ ( · ) . Finally , spatial global average pooling ρ ( · ) is applied on top of them to reduce the dimensionality to RC . The resulting tokens are gathered to form the output tensor : Zt = [ zi ] Si=1 ∈ RS×C . The overall process has a form of an element-wise spatial self-attention . In our version , { αi ( · ) } Si=1 are implemented together as a single or a series of convolutional layers ( with the channel size S ) followed by a sigmoid function , although this could be extended with other implementations . In case of an image , Z = Zt . In the case of a video , the tokens Zt from all the frames are collected to form the final output token tensor Z ∈ RST×C . We specifically name our token learning module as “ TokenLeaner ” . Figure 1 visually summarizes the TokenLearner module . Compute reduction in Transformers : The learned tokens ( i.e. , the outputs of the TokenLearner Z ) are provided to the subsequent layers for the visual representation learning , such as multi-head selfattention ( MHSA ) used in Vision Transformer and ViViT . With the TokenLearner , these subsequent layers only need to process a small number of tokens ( e.g. , 8 instead of 1024 per frame ) and this significantly reduces the computations , as they are quadratic to the number of tokens . Figure 4 ( a ) shows a basic architecture inserting the TokenLearner module within ViViT . It could be added at any location within the network , and the relative compute of the Transformer layers after the TokenLearner become almost negligible due to the huge difference in the number of tokens . 2.2 TokenFuser . After the TokenLearner generates tokens and its subsequent Transformer layer ( e.g. , MHSA ) processes them , the “ TokenFuser ” could be used to further ( 1 ) fuse information across the tokens and ( 2 ) remap the representation back to its original spatial resolution . This enables the model to capture spatial ( or spatio-temporal ) ‘ patterns ’ formulated by the tokens , and recover the original input tensor shape when necessary . First , given the token tensor Y ∈ RST×C from a Transformer layer , we apply a linear layer ( i.e. , a fully connected MLP layer ) over the tokens , not channels . That is , we learn a linear function of RST 7→ RST where S is the number of our tokens mined per frame and T is temporal size of the input tensor , and apply it to every channel independently . That is , we update Y = ( Y TM ) T where M is a learnable weight matrix with size ST × ST . The result of such operation maintains the tensor size of ST × C. We believe this also has a connection to the observations from the concurrent work , MLPMixer [ 36 ] , that token-wise linear layers are beneficial . Next , the TokenFuser processes each temporal slice Yt ∈ RS×C individually , and remaps the token tensor of size S×C back to H ×W ×C , by learning to combine the tokens for each spatial location in H ×W differently . Xj+1t = B ( Yt , X j t ) = BwYt +X j t = βi ( X j t ) Yt +X j t ( 2 ) where Xjt is the residual input to the previous TokenLearner module , Yt is the processed tokens in the TokenFuser module , and Xj+1t is the output . Bw ∈ RHW×S is an intermediate weight tensor computed with the function βi ( Xt ) . The function βi ( Xt ) is implemented with a simple linear layer followed by a sigmoid function . Figure 2 illustrates the overall process of the TokenFuser ( the token-wise linear layer is omitted ) . 2.3 Video architecture overview . Here , we provide an overview of video representation architecture with TokenLearner . The TokenLearner and TokenFuser modules introduced in Section 2 are directly applicable for video representation learning . TokenLearner generates multiple Zt for frames in videos and they are stacked to form Z . Once Z is generated , any standard Transformer layers could be used to parse them jointly . Figure 3 provides an overview of the combined architecture for video representation , which is to be repeated over multiple layers . TokenLearner first extracts S number of tokens per frame , resulting in a total of ST tokens where T is the number of frames . Once TokenLearner generates these adaptively learned tokens , they are provided to the subsequent Transformer layer to capture the global space-time patterns . Finally ( and optionally depending on the architecture ) , TokenFuser applies a linear layer over the token axis and remaps the tensor shape back , as discussed in Subsection 2.2 . Following Eq . 2 , TokenFuser is applied for per-frame representation Yt . This results in a lightweight approach , which brings forth an efficient video representation by capturing long-range visual patterns .

Summary: This paper introduces a method of adaptive space-time tokenization (TokenLearner) for video representation learning. Instead of using fixed/pre-trained tokens for videos, TokenLearner aims to generate multiple spatial attention maps for the architecture to better focus on regions with rich information, thus obtaining better tokens. The overall architecture mimics a combination of X(2+1)D and Bottleneck Transformer while inserting the TokenLearner module. The authors conduct experiments on Charades and AViD datasets, and the proposed method outperforms previous state-of-the-art models. The authors have also carried out ablation study on their design choices.