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: The paper proposes a transformer-based neural network architecture for extracting video representations. Past approaches discretize entire videos into 3D (RGB-T) chunks as input 'tokens' to transformers. Instead in this work, an additional operation is proposed to extract a smaller number of tokens per frame. This leads to a more computationally efficient architecture that outperforms existing methods on the task of video classification.