STDAN: Deformable Attention Network for Space-Time Video Super-Resolution

We first use a $3\times 3$ convolutional layer in the feature extraction module to get shallow features $\left\{F_{2t-1}^{s}\right\}_{t=1}^{N}$ from the $N$ input LR video frames. Considering that these shallow features lack long-range spatial information due to the locality of the naive convolutional layer, which may cause poor quality in the next feature interpolation module. We hope to extract these shallow features further to establish the correlation between two distant locations.

Recently, Transformer-based models have realized good performance in computer vision [5, 7, 6, 10] owing to the strong capacity of Transformer to model long-range dependency. However, the computation cost of self-attention in Transformer is high, which limits its extensive application in video-related tasks. To overcome the drawback, Liu et al. [11] put forward Swin Transformer block (STB) to achieve linear computational complexity with respect to image size. Based on efficient and effective STB [11], Liang et al. [12] proposed residual Swin Transformer block (RSTB) to construct SwinIR for image restoration. Thanks to the powerful ability to model long-range dependency of RSTB, SwinIR [12] obtains state-of-the-art (SOTA) performance compared with CNN-based methods. In this paper, to acquire features $\left\{F_{2t-1}\right\}_{t=1}^{N}$ that capture long-range spatial information, we also use $m_{f}$ RSTBs [12] to extract shallow features $\left\{F_{2t-1}^{s}\right\}_{t=1}^{N}$ further, illustrated in Fig. 2. We can see that the RSTB is a residual block with several STBs and one convolutional layer. In addition, given a tensor $X_{in}$ as input, the detailed process of STB to output $X_{out}$ is formulated as:

where MSA and $X$ denotes multi-head self-attention module and intermediate results, respectively.

To implement the super-resolution in the time dimension, we also utilize a feature interpolation module to synthesize the intermediate frames in the LR feature space, like [13, 36]. Specifically, given the two extracted features: $F_{1}$ and $F_{3}$, the feature interpolation module can synthesize the feature $F_{2}$ corresponding to the missing frame $I_{2}^{lr}$. Generally, to obtain the intermediate feature, we should capture the pixel-wise motion information first. Optical flow is usually adopted to estimate the motion between video frames. However, there are several shortcomings in using optical flow for interpolation. The computational cost is high to calculate optical flow precisely, and estimated optical flow may be inaccurate due to the occlusion or motion blur, which causes poor interpolation results.

Considering the drawback of optical flow, Xiang et al. [13] employed multi-level deformable convolution [38] to perform frame feature interpolation. The learned offset used in deformable convolution can implicitly capture forward and backward motion information and achieve good performance. However, the synthesis of intermediate frame feature [13, 36] only utilizes the two neighboring frame features, which cannot fully explore the information from the other input frames to assist in the process. Unlike feature interpolation in previous STVSR algorithms [13, 36], we propose a long-short term feature interpolation (LSTFI) module to realize the intermediate frame in our STDAN, which is capable of exploiting helpful information from more input frames.

As illustrated in Fig. 3, we adopt a bidirectional recurrent neural network (RNN) [37] to construct the LSTFI module, which consists of two branches in forward and backward directions. Take the forward branch as an example. Two neighboring frame features and the hidden state from the previous long-short term cell (LSTC) are fed into each LSTC, and then the LSTC generates the corresponding intermediate frame feature and current hidden state used for subsequent LSTC. Here, the two neighboring frame features and hidden state serve as short-term and long-term information for the intermediate feature, respectively. However, each branch’s hidden state only considers the unidirectional information flow. To fully mine the information flow of these frame features for the interpolation procedure, we fuse interpolation results from LSTCs in the forward and backward branches to acquire the final intermediate frame feature.

The architecture of LSTC and the fusion process are shown in Fig. 4. Given two neighboring frame features $F_{1}$ and $F_{3}$, we employ deformable feature interpolation (DFI) block [13] to capture the forward and backward motion between the two features implicitly. For simplification, we take the feature $F_{3\rightarrow{1}}$ that has experienced backward motion compensation as an example. As illustrated in Fig. 4(b), the two frame features are concatenated along channel dimension, and then pass through offset generation function $H_{og}^{b}$ to predict an offset with backward motion information:

where $H_{og}^{b}$ consists of convolutional layers, and $[.,.]$ denotes the concatenation along channel dimension. With the learned offset, we adopt deformable convolution [40] as motion compensation function to obtain compensated feature:

where $DConv$ denotes the operation of deformable convolution.

To blend the features $F_{1\rightarrow{3}}$ and $F_{3\rightarrow{1}}$ that have experienced forward and backward motion compensation respectively, a $1\times{1}$ convolutional layer is applied, which can perform pixel-level linear weighting to achieve preliminary interpolation feature $F_{2}^{p}$. Note that the acquisition of feature $F_{2}^{p}$ only utilizes the short-term information. In order to combine long-term information $h_{0}^{f}$, the hidden state from the previous LSTC, we first use the other DFI block to align $h_{0}^{f}$ with the current feature $F_{2}^{p}$, since there may be some misalignment. The process is expressed as:

where $DAlign(.)$ indicates the operation of DFI block. At the end of LSTC, we apply a fusion function into aligned hidden state $h_{0\rightarrow{2}}^{f}$ and preliminary interpolation result $F_{2}^{p}$ to obtain forward intermediate feature:

where $H_{fs}$ refers to the fusion function. Then, the intermediate feature $F_{2}^{f}$ passes through a convolutional layer and an activation layer in sequence to produce hidden state $h_{2}^{f}$ for the subsequent LSTC.

For fully exploring the whole input frame features for interpolation, the bidirectional RNN structure is utilized in our LSTFI module, so we fuse the forward intermediate feature $F_{2}^{f}$ and backward intermediate feature $F_{2}^{b}$ to get the final intermediate frame feature $F_{2}$, shown in Fig. 4(c).

With the assistance of the LSTFI module, we now have $2N-1$ frame features, where the generation of $N-1$ intermediate frame features combines their adjacent frame features with hidden states. Although the hidden states can introduce certain temporal information, the whole interpolation procedure hardly explicitly explores the temporal information between various frames. In addition, the $N$ input frame features are merely processed independently in the feature extraction module. However, these frame features $\left\{F_{t}\right\}_{t=1}^{2N-1}$ are consecutive, which means there are abundant temporal content without being exploited among these features.

For a feature vector $\mathbf{f}_{t}$ whose location is $\mathbf{p}_{o}$ on feature $F_{t}$, the simplest method to aggregate temporal information is adaptive fusion with the feature vector on the same location from the other $2N-2$ frame features. However, the aggregation approach has several drawbacks. Generally, the corresponding point on other frame features may not be on the same location due to inter-frame motion. Furthermore, there are multiple helpful feature vectors for $\mathbf{f}_{t}$ from each of the $2N-2$ frame features. Based on the above analysis, we propose a spatial-temporal deformable feature aggregation (STDFA) module to mix cross-frame spatial information adaptively and capture the long-range temporal information.

Specifically, we utilize the STDFA module to learn the residual auxiliary information from the remaining $2N-2$ frame features for each frame feature $F_{t}$. As presented in Fig. 5, the processing of the STDFA module can be divided into two parts: spatial aggregation and temporal aggregation. To adaptively fuse cross-frame spatial content of frame feature $F_{i}$ from the other frame features, we perform deformable attention to each pair: $F_{i}$ and $F_{j}$ ($j\in[1,2N-1]$, $j\neq{i}$). In detail, frame feature $F_{i}$ passes through a linear layer to get embedded feature $Q_{i}$. Similarly, frame feature $F_{j}$ is fed into two linear layers to obtain embedded features $K_{j}$ and $V_{j}$, respectively.

To implement deformable attention between $F_{i}$ and $F_{j}$, we first predict the offset map:

where $H_{og}$ indicates offset generation function consisting of several convolutional layers with $k\times{k}$ kernel. The offset map $\Delta{M_{j\rightarrow{i}}}$ at position $\mathbf{p}_{o}$ is expressed as:

Then the offsets $\Delta{M_{j\rightarrow{i}}(\mathbf{p}_{o})}$ are combined with $k^{2}$ pre-specified sampling locations to perform deformable sampling. Here, we denote the pre-specified sampling location as $\mathbf{p}_{\xi}$, and the value set of $\mathbf{p}_{\xi}$ of $k\times{k}$ kernel is defined as:

where $\left\lfloor\cdot\right\rfloor$ denotes rounding down function.

With the offsets $\Delta{M_{j\rightarrow{i}}}(\mathbf{p}_{o})$, the embedded feature vector $Q_{i}(\mathbf{p}_{o})$ can attend $k^{2}$ related points in $K_{j}$. Nevertheless, not all the information of these $k^{2}$ points is helpful for $Q_{i}(\mathbf{p}_{o})$. In addition, each point on embedded feature $Q_{i}$ needs to search $k^{2}$ points, which inevitably causes a large storage occupation. To avoid irrelevant points and reduce storage occupation, we only choose the first $T$ points that are most relevant. To select the $T$ points, we calculate the inner product between two embedded feature vectors as the relevance score:

The larger the score, the more relevant the two points are. According to this criterion, we can determine the $T$ points. In the following, to distinguish the selected $T$ points from original $k^{2}$ points, we denote the pre-specified sampling location and learned offset of the $T$ points as $\mathbf{\overline{p}}_{\xi}$ and $\Delta{\mathbf{\overline{p}}_{\xi}}$, respectively.

To adaptively mingle the spatial information from the $T$ locations for each embedded feature vector $Q_{i}(\mathbf{p}_{o})$, we first adopt softmax function to calculate the weight of these points:

Then, with the weights and the embedded feature vector $K_{j}(\mathbf{p}_{o}+\mathbf{\overline{p}}_{\xi}+\Delta{\mathbf{\overline{p}}_{\xi}})$, we can obtain corresponding updated embedded feature vector:

Same as $K_{j\rightarrow{i}}(\mathbf{p}_{o})$, the updated vector $V_{j\rightarrow{i}}(\mathbf{p}_{o})$ can be also achieved with the weight $w_{\xi}$. Finally, we calculate the updated relevant weight map $W_{j\rightarrow{i}}$ at each position $\mathbf{p}_{o}$ between $Q_{i}$ and $K_{j\rightarrow{i}}$ for the following temporal aggregation:

To capture the temporal contexts of frame feature vector $F_{i}(\mathbf{p}_{o})$ from the remaining $2N-2$ features, we also utilize softmax function to adaptively aggregating feature vectors $V_{j\rightarrow{i}}(\mathbf{p}_{o})$. Specifically, the normalized temporal weight of each vector $V_{j\rightarrow{i}}(\mathbf{p}_{o})$ ($j\in[1,2N-1]$, $j\neq{i}$) is expressed as:

Then, through fusing embedded feature vector $V_{j\rightarrow{i}}(\mathbf{p}_{o})$ ($j\in[1,2N-1]$, $j\neq{i}$) with the corresponding normalized weight, we can attain the embedded feature $V_{i}^{*}$ that aggregates the spatial and temporal contexts from other $2N-2$ embedded features. The weighted fusion process is defined as:

In the tail of STDFA module, the embedded feature $V_{i}^{*}$ is sent into a linear layer to acquire the residual auxiliary feature $F_{i}^{res}$. Finally, we add frame feature $F_{i}$ and residual auxiliary feature $F_{i}^{res}$ to get the enhanced feature $F_{i}^{*}$ that aggregates spatial and temporal contexts from the other $2N-2$ frame features.

To reconstruct HR frames from the enhanced features $\left\{F_{t}^{*}\right\}_{t=1}^{2N-1}$, we first employ $m_{b}$ RSTBs [12] to map feature $F_{t}^{*}$ to deep feature $F_{t}^{d}$. Then, these deep features further pass through an upsampling module to realize the HR video frames $\left\{I_{t}^{hr}\right\}_{t=1}^{2N-1}$. Specifically, the upsampling module consists of the PixelShuffle layer [16] and several convolutional layers. For optimizing our proposed network, we adopt the Charbonnier function [17] as the reconstruction loss:

where $I_{t}^{GT}$ indicates the ground truth of $t$-th reconstructed video frame $I_{t}^{hr}$, and the value of $\epsilon$ is empirically set to $1\times 10^{-3}$. With the loss function, our STDAN can be end-to-end trained to generate HR slow-motion videos from corresponding LR and LFR counterparts.

Datasets We use the Vimeo-90K dataset [32] to train our network. Specifically, the Vimeo-90K dataset consists of more than 60,000 training video sequences, and each video sequence has seven frames. We adopt the raw seven frames as our HR and HFR supervisions. The corresponding four LR and LFR frames are downscaled by a factor of 4 with bicubic sampling from these odd-numbers ones. The Vimeo-90K also provides corresponding testsets that can be divided into Vimeo-Slow, Vimeo-Medium and Vimeo-Fast according to the degree of motion. The three testsets serve as the evaluation datasets in our experiments. Same as STVSR methods [13, 36], six video sequences in Vimeo-Medium testset and three sequences in Vimeo-Slow testset are removed to avoid infinite values on PSNR. In addition, we report the results on Vid4 [33] and SPMC-11 [23] of different approaches.

In our STDAN, the number of RSTBs in the feature extraction module and frame feature reconstruction module is 2 and 6, respectively, where each RSTB contains 6 STBs. In addition, the number of feature and embedded feature channels are set to be 64 and 72 separately. In the LSTFI module, we utilize a Pyramid, Cascading, and Deformable (PCD) structure in [38] to achieve DFI. The hidden states in the forward and backward branches are initialized to zeros. In the STDFA module, the value of $k$ and $T$ are set to 3 and 2, respectively. We augment the training frames by randomly flipping horizontally and $90^{\circ}$ rotations during the training process. Then, we crop the input LR frames with a size of $32\times 32$ at random to the network, and the batch size is set to be 18. Our model is trained by Adam [35] optimizer by setting $\beta_{1}=0.9$ and $\beta_{2}=0.999$. We employ cosine annealing to decay the learning rate [41] from $2e-4$ to $1e-7$. We implement the STDAN with PyTorch and train our model on 6 NVIDIA GTX-1080Ti GPUs.

We compare our STDAN with existing state-of-the-art (SOTA) one-stage STVSR approaches: STARnet [66], Zooming Slow-Mo [13], RSTT [65] and TMNet [36]. In addition, we also compare the performance of our network with SOTA two-stage STVSR algorithms, like Zooming Slow-Mo [13] and TMNet [36]. Specifically, two-stage STVSR methods are composed of SOTA VFI and SR algorithms. These VFI networks are SuperSloMo [25], SepConv [30] and DAIN [31], respectively, while SOTA SR approaches are RCAN [56], RBPN [52] and EDVR [38].

Quantitative results of various STVSR methods are shown in Table I. From the table, we can see that: (1) Our STDAN with fewer parameters obtains SOTA performance on both Vid4 [33] and Vimeo [32]; (2) For the SPMC-11 [23] dataset, our model is only 0.1dB lower than STARnet [66] in terms of PSNR, but our STDAN acquires better results than it on SSIM [34] index, which demonstrates our network can recover more correct structures. In addition, our model only needs one thirteenth parameters of STARnet.

Visual comparison of different models are displayed in Fig. 6. We observe that our STDAN, with the proposed LSTFI and STDFA modules, restores more accurate structures and fewer motion blurs compared with other STVSR approaches, which confirms the higher value on PSNR and SSIM achieved by our model.

To investigate the effect of the proposed modules in our STDAN, we conduct comprehensive ablation studies in this section.

Quantitative results on Vid4 [33] and Vimeo-Fast [32] datasets are shown in Table II. From the table, we know that: (1) Feature aggregation module can improve the reconstruction results; (2) The larger the spatial range of feature aggregation, the more useful information can be captured to enhance recovery quality of HR frames. Qualitative results of the three models are represented in Fig. 8, which confirms the feature aggregation in the deformable window can acquire more helpful content.