Memory-Augmented Non-Local Attention for Video Super-Resolution

Fig. 2 demonstrates the structure of our video super-resolution network. The goal of our network is to super-resolve a single low-resolution frame $\mathbf{I}_{t}\in\mathbb{R}^{3\times H\times W}$, given the low-resolution temporal neighbor frames $\left\{\mathbf{I}_{t-\tau},...,\mathbf{I}_{t+\tau}\right\}$, where $H$ and $W$ are the video height and width respectively. To make the discussion more concise, we will use “current frame” to refer to $\mathbf{I}_{t}$ and “neighbor frames” to refer to $\left\{\mathbf{I}_{t-\tau},...,\mathbf{I}_{t+\tau}\right\}$. We will use $T=2\tau+1$ to represent the time span of neighbor frames. Note that neighbor frames include the current frame.

The first stage of our network embeds all the video frames into the same feature space by applying the same encoding network to each input frame. We denote the embedded features as $\left\{\mathbf{F}_{t-\tau},...,\mathbf{F}_{t+\tau}\right\}\in\mathbb{R}^{C\times H\times W}$, where $C$ is the dimension of the feature space.

As we discussed in Sec. 1, our super-resolution process refers to both the current video and general videos. Based on this principle, we adapt the attention mechanism which allows us to query the pixels that need to be super-resolved in the keys consist of auxiliary pixels. Specifically, the second stage of our network includes two parts: Cross-Frame Non-local Attention and Memory-Augmented Attention.

Cross-Frame Non-local Attention aims to mine useful information from neighbor frame features. In this module, neighbor frame features are queried by the current frame feature and the most possible match will be selected as the output. We denote the output of the cross-frame non-local attention module as $\mathbf{X}_{t}\in\mathbb{R}^{C^{\prime}\times H\times W}$, where $C^{\prime}=C/2$ is the dimension of the embedding space of the cross-frame non-local attention module. The design of this module will be discussed in Sec. 3.2.

Memory-Augmented Attention maintains a global memory bank $\mathbf{M}\in\mathbb{R}^{C^{\prime}\times N}$ to memorize useful information from general videos in the training set, where $N$ represents an arbitrary number of entries in the memory bank. We use the current frame feature to query the memory bank directly. However, unlike the cross-frame non-local attention module in which the keys are embedded versions of neighbor frame features, the memory bank is completely learned. The output of this module is denoted as $\mathbf{Y}_{t}\in\mathbb{R}^{C^{\prime}\times H\times W}$. This module will be discussed in Sec. 3.3.

Finally, the output of the cross-frame non-local attention module $\mathbf{X}_{t}$ and memory-augmented attention module $\mathbf{Y}_{t}$ are convolved by two different convolutional layers with kernel size 1 and added to the input current frame feature $\mathbf{F}_{t}$ as residuals. A decoder decodes the output of attention modules and an up-sampling module shuffles the pixels to generate a high-resolution residual. The residual adds details to the bilinearly up-sampled blurry low-resolution frame, resulting in a clear high-resolution frame.

One of the major procedures in the conventional video super-resolution methods is to align the neighbor frames so that the corresponding pixels can be fused and improve the quality of the super-resolution of the current frame. To achieve the alignment, the typical approaches in video super-resolution works include optical flow [20, 36] and deformable convolution [28, 31]. However, aligning pixels according to color consistency is known to be a challenging task under large motion or illumination change. As a consequence, the inaccuracy in alignment will negatively impact the performance of video super-resolution. In our work, we seek to avoid this performance overhead. As we discussed in Sec. 2, the non-local attention [32] enables capturing temporally and spatially long distance correspondence. Therefore, the frame alignment can be omitted if non-local attention is used to query pixels of the current frame in neighbor frames.

The cross-frame non-local attention module is demonstrated in Fig. 3. We first normalize the input frame features using group normalization [35], resulting in the normalized neighbor frame features $\left\{\mathbf{\overline{F}}_{t-\tau},...,\mathbf{\overline{F}}_{t+\tau}\right\}$. In our non-local attention setup, the center feature $\mathbf{\overline{F}}_{t}$ is used as the query tensor, and neighbor frame features $\left\{\mathbf{\overline{F}}_{t-\tau},...,\mathbf{\overline{F}}_{t+\tau}\right\}$ serve as both the key and value tensors. The embedded version of query, key and value tensor are noted as $\mathbf{Q}\in\mathbb{R}^{C^{\prime}\times H\times W}$, $\mathbf{K}\in\mathbb{R}^{C^{\prime}\times T\times H\times W}$ and $\mathbf{V}\in\mathbb{R}^{C^{\prime}\times T\times H\times W}$ in Fig. 3. In the traditional setup of non-local attention, the next step is to flatten the temporal and spatial dimension of $\mathbf{Q}$ and $\mathbf{K}$ to $\mathbf{\widehat{Q}}\in\mathbb{R}^{HW\times C^{\prime}}$ and $\mathbf{\widehat{K}}\in\mathbb{R}^{C^{\prime}\times HWT}$ and calculating the correlation matrix $\mathbf{\Gamma}=\mathbf{\widehat{Q}}\mathbf{\widehat{K}}$. However, the size of $\mathbf{\Gamma}$ is $HW\times HWT$. Even though the input to our network is low-resolution frames, the size of this matrix is large and grows quadratically when involving more neighbor frames. This makes the training difficult due to the limited GPU memory. Moreover, note that the number of columns in $\mathbf{\Gamma}$ depends on the frame size. With the same pre-trained network, normalizing $HWT$ dimensions using softmax makes the columns in $\mathbf{\widehat{K}}$ that are most correlated with a row in $\mathbf{\widehat{Q}}$ less significant when the input frame size is larger. Intuitively, traditional non-local attention weakens the contribution of the most useful auxiliary pixels in the super-resolution case. In Sec. 4.3, we will show that traditional non-local attention did not benefit the video super-resolution method PFNL [42] which directly applies it to the entire group of neighbor frames.

To mitigate the GPU memory issue, we conduct non-local attention on each neighbor frames separately. Moreover, for each pixel in the query tensor, we conduct non-local attention only on its spatial neighbor region in the temporal neighbor frames. The insight is that video motion is limited during a short time period, and a pixel’s correspondences in the neighbor frames lie in the neighborhood of its position. Specifically, we unfold each temporal slice of $\mathbf{K}$ with dimension $C^{\prime}\times H\times W$ into $9\times 9$ patches in the spatial domain. Therefore, for each pixel $\mathbf{\widehat{Q}}_{p}\in\mathbb{R}^{1\times C^{\prime}}$ in the query tensor $\mathbf{\widehat{Q}}$, the corresponding key tensor becomes $\mathbf{K}_{p}\in\mathbb{R}^{C^{\prime}\times 81}$. As a result, the correlation matrix $\mathbf{\Gamma}_{p}=\mathbf{\widehat{Q}}_{p}\mathbf{K}_{p}$ has a dimension of $1\times 81$. Note that $\mathbf{K}_{p}$ is different for each row of $\mathbf{\widehat{Q}}$; each $\mathbf{K}_{p}$ only represents the spatial neighborhood of the specific pixel in the current frame. By stacking the $\mathbf{\Gamma_{p}}$ for each pixel, we obtain the final correlation matrix $\mathbf{\Gamma}$.

To resolve the performance degradation issue in the traditional non-local attention, we reduce the non-local attention into an one-hot attention. Specifically, for each row in $\mathbf{\Gamma}$, we only select the largest entry as follows:

$$s_{i}=\max_{j}{\mathbf{\Gamma}(i,j)},\quad d_{i}=\mathop{\mathrm{argmax}}_{j}{\mathbf{\Gamma}(i,j)}$$

(1)

where $s_{i}$ is the attention score and $d_{i}$ is the column coordinate in $\mathbf{K}_{p}$. Denote the one-hot attention results as $\mathbf{S}\in\mathbb{R}^{HW\times 1}$ and $\mathbf{D}\in\mathbb{R}^{HW\times 1}$ which consists of $s_{i}$ and $d_{i}$ respectively. We then select the value tensor $\mathbf{V}_{p}$ as follows:

$$\mathbf{\widehat{X}}_{t}=\mathbf{S}\cdot\mathbf{V}_{p}(\mathbf{D})$$

(2)

In Eqn. 2, $\mathbf{V}_{p}\in\mathbb{R}^{C^{\prime}\times 81}$ is the value tensor unfolded in the same way as $\mathbf{K}_{p}$. The operator $(\ \ )$ selects $HW$ columns from $\mathbf{V}_{p}$ using the $HW$ indices in $\mathbf{D}$. Then the selected $HW$ columns are multiplied by $HW$ attention scores in $\mathbf{S}$. The $\mathbf{\widehat{X}}_{t}\in\mathbb{R}^{C^{\prime}\times HW}$ is reshaped to $\mathbf{X}_{t}\in\mathbb{R}^{C^{\prime}\times H\times W}$ as the output of the cross-frame non-local attention module.

Cross-frame non-local attention enables the fusion of the information from neighbor frames in the current video. However, the neighbor frames used in the attention are also low-resolution with similar content to the current frame. Therefore, the benefit from cross-frame non-local attention is limited. We seek to refer to more local detail information beyond the current video, which requires memorizing useful information from the entire training set. For this purpose, our network includes a memory-augmented attention module. The module maintains a global memory bank $\mathbf{M}\in\mathbb{R}^{C^{\prime}\times N}$ which is learned as parameters of the network. We use regular non-local attention to query current frame features $\mathbf{\widehat{Q}}$ in the global memory bank $\mathbf{M}$, i.e. the correlation matrix is $\mathbf{\Gamma}_{M}=\mathbf{\widehat{Q}}\mathbf{M}\in\mathbb{R}^{HW\times N}$. Finally, we obtain the output

$$\mathbf{\widehat{Y}}_{t}=softmax(\mathbf{\Gamma}_{M})\mathbf{\widehat{M}}$$

(3)

where $\mathbf{\widehat{M}}\in\mathbb{R}^{N\times C^{\prime}}$ is the transposed version of the memory bank $\mathbf{M}$. Similar to the cross-frame non-local attention module, we reshape $\mathbf{\widehat{Y}}_{t}\in\mathbb{R}^{HW\times C^{\prime}}$ to $\mathbf{Y}_{t}\in\mathbb{R}^{C^{\prime}\times H\times W}$ as the output of the memory-augmented attention module.

In the first stage, we fix the memory-augmented attention module and train the rest part of the network for 90,000 iterations at the learning rate of $10^{-4}$. The loss function used is $L_{1}=\left\|\mathbf{O}_{t}-\mathbf{G}_{t}\right\|_{1}$, where $\mathbf{O}_{t}$ stands for the output super-resolved current frame and $\mathbf{G}_{t}$ is the ground truth high-resolution frame.

In the second stage, we fix the network weights except for the memory-augmented attention module. The loss function $L_{2}=\left\|\mathbf{Y}_{t}-\mathbf{Q}\right\|_{1}$ focus on training the memory bank. Note that the training process optimizes the memory bank $\mathbf{M}$ so that a query $\mathbf{Q}$ can be represented by the combination of the columns in $\mathbf{M}$ as accurate as possible. This is essentially clustering and summarizing the most representative general pixel features in the encoded space. We train this stage for 30,000 iterations at the learning rate of $10^{-4}$.

In the final stage, we fine-tune the entire network using $L_{1}$ for 30,000 iterations at the learning rate of $10^{-5}$.

As discussed in Sec. 1, the cross-frame non-local attention in our method enables VSR without frame alignment. To validate the robustness of our method to large motion videos, we randomly collect 14 parkour video clips from the Internet. Parkour is a form of extreme sport focusing on passing obstacles in a complex environment by running, climbing, and jumping. Usually taken using egocentric wearable cameras, parkour videos are typical examples in the real-world where large camera motions are everywhere. Example video stills from the Parkour dataset are shown in Fig. 6. We further evaluate our method on regular small motion videos using Vimeo90K [36] test set and Vid4 [20]. For all the test sets, we use the average PSNR and SSIM [49] on the RGB channels to quantitatively evaluate the performance of the methods. In addition, we apply LPIPS [44] to evaluate the perceptual similarity between the super-resolved frames and the ground truth high-resolution frame. Since the performance can be different across computation platforms and the quantitative metric calculation might be different in these works, we re-run their code and calculate the metrics in the same way on the same computer.

The visual comparison of the examples from the Parkour dataset, Vimeo90K [36] dataset and Vid4 [20] is shown in Fig. 7. To make the discussion concise, we label the ID at the bottom left of each video. We also added arrows pointing at the regions we will be discussing.

Example (a), (b), (c), (d) and (e) are selected from the Parkour dataset. These examples contain large motion and are challenging to existing VSR methods. Our method can reconstruct repetitive patterns like Example (a) and (b), while explicit frame alignment methods TOFlow [36] and TGA [10] fail due to the inaccurate frame alignment. EDVR [31] result is more blurry than our result in example (a) and (b), and the blurry issue is more visible when viewed in dynamics as we will show in the supplementary video. This indicates that the deformable convolution alignment cannot handle the alignment with large frame displacement. The VSR and SISR methods PFNL [42] and CSNLN [18] using non-local attention also suffer from the blurry issue, potentially due to the non-local attention performance degradation problem discussed in Sec. 3.2.

Example (c) focuses on general details of objects. The frame-aligning VSR methods introduce ghosting artifacts (EDVR [31]) and deformation (TOFlow [36] and TGA [10]) due to the inaccurate alignment. PFNL [42] and CSNLN [18] results have less detail than ours, indicating that our one-hot non-local attention improves the quality of regular non-local attention. Example (d) focuses on human face shape and details. As shown in the bicubic result, the original facial details are completely lost due to the down-sampling. Our method reconstructs visually pleasing details of human faces thanks to the memory-augmented module, while the comparison methods introduce blur (EDVR [31], TOFlow [36] and PFNL [42]) or reconstruct shapes that do not look like a human (TGA [10] and CSNLN [18]).

Example (e) contains thin structures. Similar to examples (a) and (b), failure in frame alignment has negatively affected the VSR methods. In this case, the performance of VSR methods EDVR [31], TOFlow [36] and TGA [10] are even worse than the SISR method CSNLN [18]. Our method with one-hot non-local attention can achieve a comparable result to CSNLN [18] in this example since our network does not require frame alignment.

As we will discuss in Sec. 4.3, the overall average quantitative metric score of our method is slightly inferior to that of EDVR [31] and TGA [10] in the Vimeo90K [36] and Vid4 dataset [20] which are relatively easy for frame aligning VSR methods. However, a larger deviation to the ground truth does not always indicate worse performance. Example (f) and (g) are selected from the Vimeo90K [36] test set. Our method tends to produce visually sharper results than EDVR [31] and TGA [10], which is often more preferred in the video super-resolution task. Example (h) is a widely used example in Vid4 [20]. In this example, EDVR [31] and TGA [10] generate the best results, but our result is comparable to their results.

To further evaluate the robustness of our method in real-world scenarios, we provide additional results in the supplementary material. These videos are arbitrarily selected from different types of videos, covering animation, movies, and vlogs. We encourage the readers to read the supplementary PDF for a more complete visual comparison.

In Table 1(a), we compare the average PSNR and SSIM [49] of our method with the comparison methods on the Parkour dataset. In this table, larger PSNR and SSIM and smaller LPIPS loss indicate better results. We mark the best result in red and the second best result in blue.

The videos in the Parkour dataset have extremely large motions, making the accurate alignment of the frames difficult. Among the comparison methods, TOFlow [36] explicitly estimates the optical flow for warping neighbor frames; TGA [10] uses homography to align neighbor frames; EDVR [31] implicitly align frames using learned kernel offset for deformable convolution. The performances of these methods are even inferior to that of the SISR method CSNLN [18] in the Parkour dataset, since fusing misaligned frames often cause blurry or ghosting artifacts in the result. This indicates that the large motions have a negative effect on the performance of traditional VSR methods.

The comparison method PFNL [42] uses pair-wise non-local attention on all the pixels in the entire segment. As discussed in Sec. 3.2, the traditional non-local attention is difficult to train due to the large GPU memory consumption and the performance degradation with a large number of pixels. Although the performance of their method is better than frame alignment methods EDVR [31] and TGA [10], using traditional non-local attention on all video segments is worse than applying it only on a single frame like CSNLN [18]. Our cross-frame non-local attention mechanism significantly improves the performance of non-local attention by introducing the one-hot attention in the video super-resolution task.

In addition, we provide the quantitative comparison on the Vimeo90K dataset [36], the Vid4 dataset [20] and the SPMC dataset [27] in Table 1(b), (c) and (d) respectively. The metrics and color labels are the same as Table 1(a). For these general small motion videos datasets, we also achieve better results than the explicit optical flow alignment VSR method TOFlow [36] and the other non-local attention super-resolution methods PFNL [42] and CSNLN [18]. Note that the single image non-local attention method CSNLN [18] also outperforms video non-local attention method PFNL [42] in the Vimeo90K dataset. The PSNR and SSIM value of our method is slightly inferior to that of EDVR [31] and TGA [10] in the Vimeo90K and Vid4 datasets. EDVR [31] has a 0.97dB and 0.64dB PSNR gain to our method in Vimeo90K and Vid4 respectively. TGA [10] has a 0.32dB and 0.21dB PSNR gain to our method in Vimeo90K and Vid4 respectively. The perceptual quality measured by the LPIPS [44] are similar among EDVR [31], TGA [10] and our method, with around 0.02 differences. However, for the large motion examples in the Parkour dataset, our method has a larger PSNR gain (1.87dB and 2.34dB) and LPIPS gain (0.0659 and 0.0983) in the performance comparing to EDVR [31] and TGA [10]. Moreover, in Table. 2, we computed the optical flow for videos in the Vimeo90K [36] test set and ranked them based on the average flow magnitude. It can be observed that even though EDVR [31] and TGA [10] is slightly superior on average, our method is actually better on the large motion videos in Vimeo90K [36].

We also note that EDVR [31] is biased towards Vimeo90K [36], given the significant performance drop in the other small motion video dataset SPMC [27]. This implies that our method is more robust than the comparison methods. We also provide the quantitative evaluation on the real-world videos in the supplementary material, which also supports the superior of our method in robustness.

In Table 3, we quantitatively compare the performance of different configurations in our network. Specifically, we set the memory size $N$ of the memory-augmented attention module to 128, 256, 512 and 1024. To verify the effectiveness of the memory-augmented attention module, we also experimented with the network with cross-frame non-local attention module only (labeled as No_Mem in Table 3). Among these configurations, $N=256$ achieves the best result and is selected in the comparisons in Sec. 4.2 and Sec. 4.3. Using smaller memory ($N=128$) results in slight performance degradation. The benefits saturate when using a larger memory ($N=512$ and $N=1024$), implying that the local details of low-resolution frames can be well represented in low-dimensional space. The performance of our network degrades without the memory-augmented attention module. However, solely using the cross-frame non-local attention module, our network outperforms comparison methods in the Parkour dataset and achieves comparable performance in the Vimeo90K dataset.