Spatiotemporal Self-attention Modeling with Temporal Patch Shift for Action Recognition

The video based transformer can be built by extending the image based ViT [8]. A video clip $X\in\mathbb{R}^{F\times H\times W\times C}$ can be divided into $s\times k\times k$ non-overlapped patches. The 3D patches are flattened into vectors $\mathbf{x}^{(t,p)}\in\mathbb{R}^{3sk^{2}}$ with $t=1,\dots,T$ denoting the temporal index with $T=F/s$, and $p=1,\dots,N$ denoting the spatial index with $N=HW/k^{2}$. The video patches are then mapped to visual tokens with a linear embedding layer

Suppose that a visual transformer contains $L$ encoders, each consisting of multi-head SA, Layer-Norm (LN) and FFN. The transformer encoder could be represented as follows:

Following transformer encoders, temporal and spatial averaging (or temporal averaging only if using class tokens) can be performed to obtain a single feature, which is then fed into a linear classifier. The major computation burden of transformers comes from the SA computation. Note that when full spatiotemporal SA is applied, the complexity of attention operation is $\mathcal{O}(N^{2}T^{2})$, while spatial-only attention costs $\mathcal{O}(N^{2}T)$ in total. Next, we show how to turn a spatial-only SA operator to a spatiotemporal one with TPS.

TSM [22] uses space-invariant channel shift for temporal redundancy modeling, which is a special case of our proposed patch shift operation by shifting $\alpha$ percent of channels (with $\alpha$ percent of elements in $\mathbf{a}_{i}$ equal to 1), where $\alpha$ is a constant for all patches. In our case, we mainly explore temporal spatial mixing and shift patches in a space-variant manner, where $\mathbf{a}_{i}=\mathbf{0}$ or $\mathbf{1}$. To reduce the mixing space, shift pattern $\mathbf{p}$ is introduced, which is applied repeatedly in a sliding window manner to cover all patches. For example, $\mathbf{p}=\{0,1\}$ means shifting one patch for every two patches, therefore $\mathbf{A}=[\mathbf{0},\mathbf{1},\mathbf{0},\mathbf{1},\dots]$ in Eq. 4. In practice, 2D shift patterns are designed for video data, which will be discussed in detail in the section below.

Patch shift SA. By using the proposed patch shift operation, we can turn spatial-only SA into spatiotemporal SA. Given video patches $\mathbf{Z}\in\mathbb{R}^{D\times T\times N}$, the PatchShift function shifts the patches of each frame along the temporal dimension with pattern $\mathbf{p}$. As only part of patches are shifted in each frame, patches from different frames could be presented in the current frame, therefore, spatial-only SA naturally turns into a sparse spatiotemporal SA. After SA, patches from different frames are shifted back to their original locations. We follow [23] to add a relative position bias with an extension to 3D position. To keep the track of shifted patches, the 3D positions are shifted alongside. With PatchShift, the multi-head SA is computed as:

Patch shift patterns. As mentioned before, in order to reduce the design space, our strategy is to employ repeated shift patterns, as it can scale up to different input sizes and is easy to implement. We adopt the following pattern design principles: a) Even spatial distribution. For each pattern, we uniformly sample the patches from the same frame to ensure they are evenly distributed. b) Reasonably large temporal receptive field. The temporal receptive field is set large enough to aggregate more temporal information. c) Adequate shift percentage. A higher percentage of shifted patches could encourage more inter-frame information communication. Various spatiotemporal SA models can be implemented with different patch shift patterns. We show several instantiations of $\mathbf{p}$ in Fig. 5. The numbers represent the indices of the frames that the patches are from, where “0”, “-” and “+” indicate current, previous and next frames, respectively. Pattern (a) shifts a single patch from next frame to the center of current frame. Pattern (b) shifts patches with a “Bayer filter” like pattern from previous and next frames. Pattern (c) shifts patches with a temporal field of 9, with patch from current frame in the center and patches from previous and next 4 frames around it. For window size larger than the pattern size, we spatially repeat the pattern to cover all patches. We use cyclic padding in [23] for patches that exceed the temporal boundary. In the experiment section, we will discuss the design of shift patterns by extensive ablation studies.

Patch shift is an efficient spatiotemporal SA modeling method for transformers as it only costs $\mathcal{O}(N^{2}T)$ complexity in both computation and memory, which is much less than “Joint” space-temporal SA, where $N$ and $T$ are the spatial and temporal dimension of patches. Patch shift is also more efficient than other factorized attention methods such as “Divide” [2, 1] (apply spatial-only and then temporal-only SA) and “Sparse/Local” [2, 24] (subsample in space or temporal dimension). The complexity comparison of different SA models are in Table 2.

Discussions on patch and channel shifts. Patch shift and channel shift are two zero parameter and low-cost temporal modeling methods. Patch shift is space-wise sparse and channel-wise dense, while channel shift is opposite. We show an example by applying patch shift and channel shift on three consecutive frames in Fig. 3. Here the shift operations are applied directly on RGB images for visualization, while in practice we apply shift operations on feature maps. In the output image of patch shifting, $1/4$, $1/2$ and $1/4$ patches are from frames $t_{0}-1$, $t_{0}$ and $t_{0}+1$, respectively. For channel shift, the red, green and blue colors represent frames $t_{0}-1$, $t_{0}$ and $t_{0}+1$, respectively. The output of shift operation for frame $t_{0}$ is represented as $t_{0}^{\prime}$ in the last column. As we can see from the figure, both patch shift and channel shift can capture the motion of action. Patch shift is spatially sparse while keeping the global channel information for each patch. In contrast, channel shift uses partial channel information in exchange for temporal information from other frames. Previous studies on vision transformer [5] have shown that feature channels encode activation of different patterns or objects. Therefore, replacing partial feature channels of the current frame with other frames could potentially lose important information of patterns/objects of interest. In comparison, patch shift contains full information of channels of patches. When patch shift is employed for SA modeling, it builds a sparse spatiotemporal SA with 3D relations among patches. In comparison, channel shift can be viewed as a “mix-patch” operation, by which temporal information is fused in each patch with shared 2D SA weights. Patch shift and channel shift perform shifting operations in orthogonal directions and they are complementary in nature.

Based on the proposed TPS, we can build Patch Shift Transformers (PST) for efficient and effective spatiotemporal feature learning. A PST can be built by inserting TPS into the SA module of off-the-shelf 2D transformer blocks. Therefore, our model could directly benefit from the pre-trained models on 2D recognition tasks. The details of Temporal Patch Shift blocks (TPS block for short) can be seen in Fig. 4. TPS turns spatial-only SA to spatiotemporal SA by aggregating information of patches from other temporal frames. However, it gathers information in a sparse manner and sacrifices SA within frames. To alleviate this problem, we insert one TPS block for every two SA modules (alternative shift in short) so that spatial-only SA and spatiotemporal SA could work in turns to approximate full spatiotemporal SA.

We further improve the temporal modeling ability of spatial-only SA with channel shift. Specifically, partial channels of each patch are replaced with those from previous or next frames. We call this block a temporal channel shift (TCS) block. The final PST consists of both TPS and TCS blocks. We will also implement a channel-only PST for comparison to unveil the benefits of patch shift.

Something-Something V1$\&$V2 [15] are large collections of video clips, containing daily actions interacting with common objects. They focus on object motion without differentiating manipulated objects. V1 includes 108,499 video clips, while V2 includes 220,847 video clips. Both V1 and V2 have 174 classes. Kinetics400 [17] is a large-scale dataset in action recognition, which contains 400 human action classes, with at least 400 video clips for each class. Each clip is collected from a YouTube video and then trimmed to around 10s. Diving-48 V2 [21] is a fine-grained video dataset of competitive diving, consisting of 18k trimmed video clips of 48 unambiguous dive sequences. This dataset is a challenging task for modern action recognition systems as it reduces background bias and requires modeling of long-term temporal dynamics. We use the manually cleaned V2 annotations of this dataset, which contains 15,027 videos for training and 1,970 videos for testing.

Models. We choose swin transformer [23] as our backbone network and develop PST-T, PST-B with an increase in model size and FLOPs based on swin transformer Tiny and Base backbones, respectively. In ablation study, we use Swin-Tiny considering its good trade-off between performance and efficiency. We adopt 32 frames as input and the tubelet embedding strategy in ViViT [1] with patch size $2\times 4\times 4$ by default. As PST-T and PST-B are efficient models, when comparing with SOTA methods, we introduce PST-T${\dagger}$ and PST-B${\dagger}$, which doubles the temporal attention window to 2 with slightly increased computation.

Training. For all the datasets, we first resize the short side of raw images to $256$ and then apply center cropping of $224\times 224$. During training, we follow [24] and use random flip, AutoAugment [7] for augmentation. We utilize AdamW [27] with the cosine learning rate schedule for network training. For PST-T, the base learning rate, warmup epoch, total epoch, stochastic depth rate, weight decay, batchsize are set to $10^{-3}$, 2.5, 30, 0.1, 0.02, 64 respectively. For larger model PST-B, learning rate, drop path rate and weight decay are set to $3\times 10^{-4}$, 0.2, 0.05, respectively.

Testing. For fair comparison, we follow the testing strategy in previous state-of-the-art methods. We report the results of two different sampling strategies. On Something-something V1$\&$V2 and Diving-48 V2, uniform sampling and center-crop (or three-crop) testing are adopted. On Kinetics400, we adopt the dense sampling strategy as in [1] with 4 view, three-crop testing.

To investigate the design of patch shift patterns and the use of TPS blocks, we conduct a series of experiments in this section. All the experiments are conducted on Something-something V1 with Swin-Tiny as backbone (IN-1K pretrained). For experiments on design of patch shift we use patch-only PST for clarity.

The number and distribution of shifted patches. We start with a simple experiment by shifting only one center patch along temporal dimension within each window. It can be seen from Table 2(a) that this simple shift pattern (center-one) brings significant improvements ($4.7\%$ on top-1) over the model without shifting operation (none). We then increase the number of shifting patches to $1/2$ of the total patches, however, in an uneven distribution (shift only the left half of patches). This uneven shift pattern does not improve over the simple “center-one” pattern. However, when the shifted patches are distributed evenly within the window (even-2), the performance increases by $0.9\%$ on top-1. This indicates that, the shifted and non-shifted patches should be distributed evenly. It is also found that a large temporal field is helpful. When we increase the temporal field to 3 by shifting 1/4 patches to previous frame and 1/4 patches to next frame (even-3), the performance is improved by $2.4\%$ on top-1 over shifting patches in one dimension.

Patch shift patterns. Based on the experiments in Table 2(a), we design a few different patch shift patterns with various temporal fields. The patches of different frames are distributed evenly within the window. Pattern A with temporal field 3 is shown in Fig. 5(b), which is “Bayer filter” like. Pattern B with temporal field 4 is implemented by replacing $1/4$ of index “0” frame patches in pattern A as index “2” frame patches. Pattern C is shown in Fig. 5(c). Pattern D is designed by placing patches from 16 consecutive frames in a $4\times 4$ pattern grid. More details can be found in the supplementary materials. We can see from the Table 2(b) that the performance of TPS gradually increases when the temporal field grows. The best performance is reached at temporal field 9, which achieves a good balance between spatial and temporal dimension. We use this pattern for the rest of our experiments.

The number of stages using TPS blocks. As can be seen in Table 2(c), by using TPS blocks in more stages of the network, the performance gradually increases. The model achieves the best performance when all the stages are equipped with TPS blocks.

Shift back, alternative shift and shift RPE. Shift back operation recovers the patches’ locations and keeps the frame structure complete. Alternative shift is described in Section 3.3 for building connections between patches. Shift RPE represents whether relative positions are shifted alongside patches. As shown in Table 2(d), removing each of them would decrease performance by $4.5\%,5.4\%,5.7\%$, respectively. Therefore, we use all the operations in our model.

Comparison with other temporal modeling methods. In Table 2(e), we compare PST with other designs of spatiotemporal attention methods in FLOPs, peak training memory consumption and accuracy. “Avgpool” achieves only $40.6\%$ Top-1 rate since it cannot distinguish temporal ordering. Joint spatiotemporal SA achieves a good performance at the price of high computation and memory cost. Local spatiotemporal attention [24] applies SA within a local window with 3D window size. It reduces the computation and memory cost but at the price of performance drop comparing to “Joint”. We also implement a sparse spatiotemporal SA by subsampling spatial patches to half, while keeping the full temporal dimension. It performs poorly as only part of patches participate in SA computation in each layer. We implement channel-only PST with shift ratio equals $1/4$, which is the same in the [22]. This channel-only PST also achieves strong performance.

Patch-only PST outperforms all other temporal modeling methods without additional parameters and FLOPs. The best performance comes from combining channel-only and patch-only PST, which indicates they are complementary in nature. Specifically, PST exceeds joint spatiotemporal SA with much less computation and only 1/5 of memory usage. It also outperforms other efficient spatiotemporal SA models with fewer computation and memory cost.

Something V1 $\&$ V2. The performance statistics on Something V1 $\&$ V2, including the pretrained dataset, classification results, inference protocols, the corresponding FLOPs and parameter numbers are shown in Table 3.

The first compartment contains methods based on 3D CNNs or factorized (2+1)D CNNs. Using efficient inference protocol (16 views and center crop$\times$1 clip), ImageNet1K pretrain, PST-T obtains 52.2$\%$ accuracy on V1 and 65.7$\%$ on V2, respectively, which outperforms all the CNN-based methods with similar FLOPs. Note that TDN [38] uses a different sampling strategy comparing to other methods, which requires 5 times more frames. Even though, PST-T still outperforms TDN on larger Something-something V2.

The second compartment contains transformer-based methods. Our small model PST-T pretrained on Kinetics400 offers competitive performance on SomethingV2 comparing to these methods. PST-B is a larger model and pretrained on ImageNet21K/Kinetics400. PST-B achieves $57.7\%$ and $69.2\%$ on V1 $\&$ V2, outperforming Timesformer-HR [2], ViViT-L [1] and MViT-B [9] at a lower cost on computation or parameter. PST-B also outperforms Mformer-L [29] and X-ViT [3], which are recently proposed efficient temporal modeling methods. They apply trajectory self-attention and local spatiotemporal self-attention, respectively. Our PST-B${\dagger}$ achieves $58.3\%$ and $69.8\%$ on V1 $\&$ V2, which outperforms other transformer-based methods. Note that, SIFAR-L [10] uses larger backbone network Swin-L, however, its performance is less satisfactory. PST-B${\dagger}$ also outperforms Video-Swin [24], which uses full temporal window on this dataset. The performances of PST family on Something-something V1$\&$V2 confirms its remarkable ability for spatiotemporal modeling.

Kinetics400. We report our results on scene-focused Kinetics400 in Table 4 and compare them with previous state-of-the-arts. As we can see from Table 4, PST-T achieves $78.2\%$ top-1 accuracy and outperforms majority of (2+1D) CNN-based methods such as TSM [22], TEINet [26], TEA [20], TDN [38] with less total FLOPs. Our larger model PST-B achieves $81.8\%$ top-1 accuracy, which outperforms strong 3D-CNN counterparts such as SlowFast [13] and X3D [12].

Comparing to transformer-based methods such as Timesformer [2] and ViViT-L [1], our PST-B${\dagger}$ achieves $82.5\%$ with less computation overheads. Specifically, PST-B${\dagger}$ outperforms SIFAR-L [10] which adopts larger size Swin-L as backbone network. PST-B${\dagger}$ also achieves on par performance with recently developed Video-Swin [25] with less computation cost.

Diving-48 V2. In Table 5, we further evaluate our method on fine-grained action dataset Diving-48. As older version of diving-48 has labeling errors, we compare our methods with reproduced SlowFast [13] and Timesformer in [2] on V2. PST-T achieves competitive $79.2\%$ top-1 performance with ImageNet-1K pretrain and costs much less computation resources than SlowFast, Timersformer and Timesformer-HR. When PST-T is pretrained on Kinetics400, it also outperforms Timersformer-L. Our larger model PST-B outperforms TimsFormer-L by a large margin of $2.6\%$ when both are pretrained on ImageNet-21K with 9.7$\times$ less FLOPs. When model is pretrained on Kinetic400, PST-B achieves $85.0\%$ top-1 accuracy. PST-T${\dagger}$ and PST-B${\dagger}$ further boost the performance to $82.1\%$ and $86.0\%$, respectively.

Latency, throughput and memory. The inference latency, throughput (videos/second) and inference memory consumption are important for action recognition in practice. We performed measurement on a single NVIDIA Tesla V100 GPU. We use batch size of 1 for latency measurement and batch size of 4 for throughput measurement. The accuracy is tested on validation set of Something-something V1&V2 dataset. The results are shown in Table 6. We make three observations: (1) Comparing to baseline “2D Swin-T” that uses 2D transformer with avgpooling for temporal modeling, PST-T has slightly higher latency, same memory consumption, but much better performance ($11.6\%$, $9.0\%$ top-1 improvements on Sthv1&v2, respectively). (2) Comparing to “Video-Swin-T” that utilizes full scale temporal information, PST-T uses $76\%$ less memory and has 2 times faster inference speed, while achieves very competitive performance. (3) For larger model PST-B${\dagger}$, it achieves better performance than Video-Swin-B while using $50\%$ less inference memory and $80\%$ less time.