VSR-Transformer | 超越BasicVSR，Transformer拿下视频超分

大家好，我是Happy。

最近，Transformer在CV领域搅翻了天，在各个CV领域打出了半片天。比如，Detection、Segmentation、Classification等等均出现了Transformer这个“搅局者”；甚至Low-level领域也被Transformer占领多个高峰，比如超分、降噪、去雨等等。一直以为视频超分可能还会再等一段时间，哪成想，视频超分领域也成了Transformer的附属地。

标题&作者团队

paper: https://arxiv.org/abs/2106.06847

code: https://github.com/caojiezhang/VSR-Transformer

Abstract

视频超分旨在由低分辨率视频进行高分辨率视频重建，是一种空-时序列预测问题。近来，Transformer由于其对序列到序列问题的并行计算能力获得了极大关注。乍一看，好像可以将Transformer直接应用到视频超分中。然而，Transformer的经典设计方式(全连接自注意力层与词级前馈层)并不适用于视频超分问题，主要由以下两个原因：

因其依赖于线性层计算注意力图，全连接自注意力层难以探索数据局部特性；
词级前馈层缺乏特征对齐能力，而这对于视频超分非常重要。

本文首次采用Transformer对视频超分任务进行了适配。具体来说，为解决第一个问题，我们提出一种空时卷积自注意力层以探索局部信息；为解决第二个问题，我们设计了一种双向光流前馈层以挖掘近邻视频帧相关性并进行对齐。最后，我们在几个基准数据集上验证了所提方案的有效性。

Method

上图给出了本文所提VSR-Transformer架构示意图，它包含特征提取、Transformer编码器以及重建网络。具体来说，给定视频序列，我们首先采用堆叠残差模块从视频中提取特征；然后VSR-Transformer将特征编码为连续表达序列；最后重建网络从上述表达中进行高分辨率视频重建。

Spatial-Temporal Convolutional Self-Attention

上图给出了本文所提STCSA(Spatial-Temporal Convolutional Self-Attention, STCSA)模块示意图，假设输入视频帧特征为

mathcal{X} in R^{Ttimes C times H times W}

，我们采用三个独立网络

mathcal{W}_q, mathcal{W}_k, mathcal{W}_v

提取每一帧的空域信息，其卷积核信息为

k=3,s=2,p=1

。不同于ViT采用线性投影方式从图像中提取多个块，我们参考COLA-Net采用unfolde操作按照步长s为Transformer提取

W_p times H_p

的块。此时，我们可以得到三组3D块，每组块数为

N=TWH/(W_p H_p)

。然后，我们通过如下方式生成

mathcal{Q,K,V} in R^{TN times C times W_p times H_p}

mathcal{Q} = k_1(mathcal{W}_q * mathcal{X}) \ mathcal{K} = k_1(mathcal{W}_k * mathcal{X}) \ mathcal{V} = k_1(mathcal{W}_v * mathcal{X})

其中，

k_1(cdot)

表示unfold操作。接下来，我们需要将每个query与key块reshape为

dtimes N

：

mathbf{Q} = tau(mathcal{Q}) \ mathbf{K} = tau(mathcal{K})

其中，

tau(cdot)

表示unfolde与reshape的复合操作。最后，我们计算相似矩阵

sigma_1(mathbf{Q}^Tmathbf{K})

并与

mathbf{V}

集成得到新的特征。通过上述方式，我们可以同时进行空时信息集成。最后，我们采用fold操作

k_2(cdot)

将前述所得特征重聚为尺寸为

Ctimes T times H times W

的特征并通过

mathcal{W}_o

处理得到输出特征。，总而言之，上述STCSA可以通过如下公式描述：

其中，

k_1^T(cdot)

表示reshape操作与folde操作的复合。在实验中，采用单头即可取得非常好的性能。

Spatial-temporal Positional Encoding 本文所提VSR-Transformer具有置换不变性，而VSR任务需要精确的空时位置信息。为解决该问题，我们提出采用3D固定位置编码并将其与注意力层的输入相加。具体来说，位置编码包含两个空域位置信息与一个时序位置信息。表示如下：

PE(pos, i) = begin{cases} sin(pos cdot alpha_k), text{ for } i = 2k \ cos(pos cdot alpha_k), text{ for } i= 2k 1 end{cases}

其中，

alpha_k = 1/10000^{2k/frac{d}{3}}

，k表示

[0,d/6)

之间的整数，pos表示对应维的位置信息，d表示通道维度。注：维度d应当被3整除。

Bidirectional Optical Flow-based Feed-forward

传统Transformer的全连接前馈层包含两个线性层外加GeLU激活，而这种方式在视频超分任务中无法有效探索不同帧间的相关性，进而导致较差的性能。

为解决该问题，我们提出了双向光流前馈层，见上图。具体来说，给定STCSA的输出特征

mathcal{X}

，我们首先学习近邻帧间的双向光流

overleftarrow{mathcal{O}}, overrightarrow{mathcal{O}}

;然后我们可以通过仿射函数

omega(cdot,cdot)

得到反向特征与前向特征:

overleftarrow{mathcal{X}} = omega(mathcal{X}, overleftarrow{mathcal{O}}) \ overrightarrow{mathcal{X}} = omega(mathcal{X}, overrightarrow{mathcal{O}})

实际上，我们采用SpyNet估计双向光流，即

我们在前向与反向传播阶段分别计算光流，见前面Figure3-right。然后，我们对仿射帧与视频帧进行集成以保持视频信息。为学习近邻帧相关性，我们提出采用卷积前向与反向网络，并修改全连接前馈网络如下：

实际上，我们还可以扩展两层网络到多层，所以上述公式可以重写如下：

f_2(mathcal{X}) = phi(f_1(mathcal{X}) rho(R_1(mathcal{V}, overleftarrow{mathcal{X}}) R_2(mathcal{V}, overrightarrow{mathcal{X}})))

相比ViT，所提方案可以有效捕获不同帧之间的相关行；不同于BasicVSR的递归计算光流与特征，所提方案可以有效避免递归，采用并行计算。

Experiments

训练数据：REDS与Vimeo90K；度量准则：PSNR与SSIM。

Results on REDS

上图给出了REDS数据集上不同方案的性能对比，从中可以看到：

相比其他方案，所提方案具有更高的PSNR指标；
当仅采用5帧进行训练时，BasicVSR与IconVSR比EDVR性能还要差，这意味着两者严重依赖于序列信息集成；
所提方法仅需64个通道即可取得比128通道的EDVR-L更好的性能；
所提方法可以重建更多的细节、更锐利的边缘。

Results on Vimeo90K

上图给出了Vimeo90K-T与Vid4数据集上的性能对比，从中可以看到：

相比其他方案，所提方案取得了最佳的PSNR指标；
在Vid4测试集上，所提方案的泛化性能要优于EDVR与BasicVSR，但弱于IconVSR。一个可能的原因：IconVSR采用全部测试帧，而所提方案仅采用了7帧。更重要的是：Vimeo90K-T与Vid4之间存在分布偏差。
视觉效果见下图，可以看到：所提方案具有更锐利而逼真的HR结果。