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: The paper proposes a new component for transformer architectures to model trajectory attention by modeling temporal correspondences in dyanmic scenes. They also present approximation techniques to address computation and memory challenges for high-resolution videos. The proposed contributions are demonstrated on video action recognition tasks with SOTA results on multiple datasets.