Keeping Your Eye on the Ball: Trajectory Attention in Video Transformers

1 Introduction . Transformers [ 75 ] have become a popular architecture across NLP [ 32 ] , vision [ 20 ] and speech [ 5 ] . The self-attention mechanism in the transformer works well for different types of data and across domains . However , its generic nature and its lack of inductive biases also mean that transformers typically require extremely large amounts of data for training [ 56 , 9 ] , or aggressive domain-specific augmentations [ 71 ] . This is particularly true for video data , for which transformers are also applicable [ 50 ] , but where statistical inefficiencies are exacerbated . While videos carry rich temporal information , they can also contain redundant spatial information from neighboring frames . Vanilla self-attention applied to videos compares pairs of image patches extracted at all possible spatial locations and frames . This can lead it to focus on the redundant spatial information rather than the temporal information , as we show by comparing normalization strategies in our experiments . We therefore contribute a variant of self-attention , called trajectory attention , which is better able to characterize the temporal information contained in videos . For the analysis of still images , ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . spatial locality is perhaps the most important inductive bias , motivating the design of convolutional networks [ 42 ] and the use of spatial encodings in vision transformers [ 20 ] . This is a direct consequence of the local structure of the physical world : points that belong to the same 3D object tend to project to pixels that are close to each other in the image . By studying the correlation of nearby pixels , we can thus learn about the objects . Videos are similar , except that 3D points move over time , thus projecting on different parts of the image along certain 2D trajectories . Existing video transformer methods [ 8 , 3 , 50 ] disregard these trajectories , pooling information over the entire 3D space-time feature volume [ 3 , 50 ] , or pooling axially across the temporal dimension [ 8 ] . We contend that pooling along motion trajectories would provide a more natural inductive bias for video data , allowing the network to aggregate information from multiple views of the same object or region , to reason about how the object or region is moving ( for example , the linear and angular velocities ) , and to be invariant to camera motion . We leverage attention itself as a mechanism to find these trajectories . This is inspired by methods such as RAFT [ 70 ] , which showed that excellent estimates of optical flow can be obtained from the correlation volume obtained by comparing local features across space and time . We observe that the joint attention mechanism for video transformers computes such a correlation volume as an intermediate result . However , subsequent processing collapses the volume without consideration for its particular structure . In this work , we seek instead to use the correlation volume to guide the network to pool information along motion paths . We also note that visual transformers operate on image patches which , differently from individual pixels , can not be assumed to correspond to individual 3D points and thus to move along simple 1D trajectories . For example , in Figure 1 , depicting the action ‘ kicking soccer ball ’ , the ball spans up to four patches , depending on the specific video frame . Furthermore , these patches contain a mix of foreground ( the ball ) and background objects , thus at least two distinct motions . Fortunately , we are not forced to select a single putative motion : the attention mechanism allows us to assemble a motion feature from all relevant ‘ ball regions ’ . Inspired by Nyströmformer [ 85 ] , we also propose a principled approximation to self-attention , Orthoformer . Our approximation sets state-of-the-art performance on the recent Long Range Arena ( LRA ) benchmark [ 69 ] for evaluating efficient attention approximations and generalizes beyond the video domain to long text and high resolution images , with lower FLOPS and memory requirements compared to alternatives , Nyströmformer and Performer [ 15 ] . Combining our approximation with trajectory attention allows us to significantly improve its computational and memory efficiency . With our contributions , we set state-of-the-art results on four video action recognition benchmarks . 2 Related Work . Video representations and 3D-CNNs . Hand-crafted features were originally used to convert video data into a representation amenable to analysis by a shallow linear model . Such representations include SIFT-3D [ 60 ] , HOG3D [ 39 ] , and IDT [ 76 ] . Since the breakthrough of AlexNet [ 40 ] on the ImageNet classification benchmark [ 58 ] , which demonstrated the empirical benefits of deep neural networks to learn representations end-to-end , there have been many attempts to do the same for video . Architectures with 3D convolutions—3D-CNNs—were originally proposed to learn deep video representations [ 72 ] . Since then , improvements to this paradigm include the use of ImageNet-inflated weights [ 11 ] , the space-time decomposition of 3D convolutions [ 54 , 74 , 84 ] , channel-separated convolutions [ 73 ] , non-local blocks [ 79 ] , and attention layers [ 13 ] . Optical flow-based pooling can be used instead of temporal convolutions to improve the representation ’ s robustness to camera and object motions [ 1 ] . Our approach shares this motivation . Vision transformers . The transformer architecture [ 75 ] , originally proposed for natural language processing , has recently gained traction in the computer vision domain . The vision transformer ( ViT ) [ 20 ] decomposes an image into a sequence of 16× 16 words and uses a multi-layer transformer to perform image classification . To improve ViT ’ s data efficiency , DeiT [ 71 ] used distillation from a strong teacher model and aggressive data augmentation . Transformers have also been used in a variety of vision image tasks , such as image representation learning [ 12 , 82 , 18 , 59 ] , image generation [ 51 ] , object detection [ 48 , 10 ] , video question-answering [ 35 ] , few-shot learning [ 19 ] , and image–text [ 49 , 63 , 67 , 44 , 68 ] , video-text [ 65 , 64 , 89 , 25 , 53 , 2 , 6 ] , and video-audio [ 43 , 52 , 29 ] representation learning . Attention for video recognition . The self-attention operation proposed in the transformer [ 75 ] have been adapted to video recognition tasks . Wang et al . [ 79 ] propose the non-local mean operation for video action recognition , which is equivalent to the standard transformer self-attention applied uniformly across space and time , while our proposed trajectory attention does not treat the space and time dimensions equivalently . Zhao et al . [ 87 ] propose a CNN architecture that explicitly predicts trajectories and aggregates information along them using a convolution operation . In contrast , our transformer architecture does not explicitly predict trajectories , but instead provides an inductive bias that encourages the network to consider motion trajectories where useful . Concurrent works [ 8 , 3 , 50 , 21 ] have also adapted the self-attention operation to the spatio-temporal nature of videos , however , these approaches do not have a mechanism for reasoning about motion paths , treating time as just another dimension , unlike our approach . Efficient attention . Due to the quadratic complexity of self-attention , there has been a significant amount of research on how to reduce its computational complexity with respect to time and memory use . Sparse attention mechanisms [ 14 ] were used to reduce self-attention complexity to O ( n √ n ) , and locality-sensitivity hashing was used by Reformer [ 38 ] to further reduce this toO ( n log n ) . More recently , linear attention mechanisms have been introduced , namely Longformer [ 7 ] , Linformer [ 78 ] , Performer [ 15 ] and Nyströmformer [ 85 ] . The Long Range Arena benchmark [ 69 ] was recently introduced to compare these different attention mechanisms . Temporal correspondences and optical flow . There are many approaches that aim to establish explicit correspondences between video frames as a way to reason about camera and object motion . For short-range correspondences across time , optical flow algorithms [ 30 , 66 , 70 ] are highly effective . In particular , RAFT [ 70 ] showed the effectiveness of an all-pairs inter-frame correlation volume as an encoding , which is essentially an attention map . All-pairs intra-frame correlations were subsequently shown to help resolve correspondence ambiguities [ 34 ] . For longer-range correspondences , object tracking by repeated detection [ 57 ] and data association can be used . In contrast to these approaches , our work does not explicitly establish temporal correspondences , but facilitates implicit correspondence learning via trajectory attention . Jabri et al . [ 31 ] estimate correspondences in a similar way , framing the problem as a contrastive random walk on a graph and apply explicit guidance via a cycle consistency loss . Incorporating such guidance into a video transformer is an interesting direction . 3 Trajectory Attention for Video Data . Our goal is to modify the attention mechanism in transformers to better capture the information contained in videos . Consider an input video I ∈ RT ′×3×H×W consisting of T ′ frames of resolution H ×W . As in existing video transformer models [ 8 , 3 ] , we pre-process the video into a sequence of ST tokens xst ∈ RD , for a spatial resolution of S and a temporal resolution of T . We use a cuboid embedding [ 3 , 21 ] , where disjoint spatio-temporal cubes from the input volume are linearly projected to RD ( equivalent to a 3D convolution with downsampling ) . We also test an embedding of disjoint image patches [ 20 ] . A learnable positional encoding e ∈ RD is added to the video embeddings for spatial and temporal dimensions separately , resulting in the code zst = xst+ess+e t t. Finally , a learnable classification token zcls is added to the sequence of tokens , like in the BERT Transformer [ 32 ] , to reason about the video as a whole . For clarity , we elide the classification token from our treatment in the sequel . We now have a set of tokens that form the input to a sequence of transformer layers that , as in ViT [ 20 ] , consist of Layer Norm ( LN ) operations [ 4 ] , multi-head attention ( MHA ) [ 75 ] , residual connections [ 27 ] , and a feed-forward network ( MLP ) : y = MHA ( LN ( z ) ) + z ; z′ = MLP ( LN ( y ) ) + y . ( 1 ) In the next section , we shall focus on a single head of the attention operation , and demonstrate how self-attention can realize a suitable inductive bias for video data . For clarity of exposition , we abuse the notation slightly , neglecting the layer norm operation and using the same dimensions for single-head attention as for multi-head attention . 3.1 Video self-attention . The self-attention operation begins by forming a set of query-key-value vectors qst , kst , vst ∈ RD , one for each space-time location st in the video . These are computed as linear projections of the input zst , that is , qst = Wqzst , kst = Wkzst , and vst = Wvzst , for projection matrices Wi ∈ RD×D . A direct application of attention across space-time ( called joint space-time attention [ 8 , 3 ] ) computes : yst = ∑ s′t′ vs′t′ · exp〈qst , ks′t′〉∑ s̄t̄ exp〈qst , ks̄t̄〉 . ( 2 ) In this way , each query qst is compared to all keys ks′t′ using dot products , the results are normalized using the softmax operator , and the weights thus obtained are used to average the values corresponding to the keys . Compared to a standard transformer , we have omitted for brevity the softmax temperature parameter D1/2 and instead assume that the queries and keys have been divided by D1/4 . One issue with this formulation is that it has quadratic complexity in both space and time , i.e. , O ( S2T 2 ) . An alternative is to restrict attention to either space or time ( called divided space-time attention ) : yst = ∑ s′ vs′t · exp〈qst , ks′t〉∑ s̄ exp〈qst , ks̄t〉 ( space ) ; yst = ∑ t′ vst′ · exp〈qst , kst′〉∑ t̄ exp〈qst , kst̄〉 ( time ) . ( 3 ) This reduces the complexity to O ( S2T ) and O ( ST 2 ) , respectively , but only allows the model to analyse time and space independently . This is usually addressed by interleaving [ 8 ] or stacking [ 3 ] the two attention modules in a sequence . Different to both of these approaches , we perform attention along trajectories , the probabilistic path of a token between frames.2 For each space-time location st ( the trajectory ‘ reference point ’ ) and corresponding query qst , we construct a set of trajectory tokens ỹstt′ , representing the pooled information weighted by the trajectory probability . The trajectory extends for the duration of the video sequence and its tokens ỹstt′ ∈ RD at different times t′ are given by : ỹstt′ = ∑ s′ vs′t′ · exp〈qst , ks′t′〉∑ s̄ exp〈qst , ks̄t′〉 . ( 4 ) Note that the attention in this formula is applied spatially ( index s ) and independently for each frame . Intuitively , this pooling operation implicitly seeks the location of the trajectory at time t′ by comparing the trajectory query qst to the keys ks′t′ at time t′ . Once the trajectories are computed , we further pool them across time to reason about intra-frame information/connections . To do so , the trajectory tokens are projected to a new set of queries , keys and values as usual : q̃st = W̃q ỹstt , k̃stt′ = W̃k ỹstt′ , ṽstt′ = W̃v ỹstt′ . ( 5 ) Like qst before , the updated reference query q̃st corresponds to the trajectory reference point st and contains information spatially-pooled from across the reference frame t. This new query is used to pool across the new time ( trajectory ) dimension by applying 1D attention : yst = ∑ t′ ṽstt′ · exp〈q̃st , k̃stt′〉∑ t̄ exp〈q̃st , k̃stt̄〉 . ( 6 ) Like joint space-time attention , our approach has quadratic complexity in both space and time , O ( S2T 2 ) , so has no computational advantage and is in fact slower than divided space-time attention . However , we demonstrate better accuracy than both joint and divided space-time attention mechanisms . We also provide fast approximations in Section 3.2 . A flowchart of the full trajectory attention operation is shown in tensor form in Figure 2 .

Summary: This work proposed a novel drop-in block for video transformers. Different from transformers that work on image or language modality, this block aggregates information along motion paths. The key idea is to design an efficient algorithm to approximation the path (trajectory attention), which requires quadratic dependence of computation and memory on the input size. As a solution, the proposed Motionformer model uses the (1) ViT image transformer model as the base architecture, (2) the separate space and time positional encodings of TimeSformer, (3) and the cubic image tokenization strategy as in ViViT. Experimental results on four video action recognition benchmarks show Motionformer can achieve state-of-the-art performance.