超分之TTVSR

这篇文章是2022年的CVPR，不同于之前VSR的对齐方式(即flow-based或者DCN-based)，本文提出的Trajectory-Aware Transformer for Video Super-Resolution(TTVSR)借用光流来做预对齐，主要利用Transformer的自注意力机制在运动轨迹中寻找相似度最高的patch作为当前帧的对齐帧。这种对齐充分利用了视频超分在时间信息上的相关性，其优势在于其更适用于长序列输入帧情况且大大减少了Transformer的高计算消耗。

参考文档
①源码
②REDS数据集

Abstract

1. VSR的关键在于如何利用好时间信息。之前的sliding-window型的VSR方法，比如EDVR等每个batch只使用固定的时间窗口大小，这样子利用的信息也是局限的。故本文提出的TTVSR采用BasicVSR系列(BasicVSR、BasicVSR++)的方式——采用全序列输入的方式，使得对齐可以利用全局信息。但不同的是，前者使用双向循环递归机制来利用全局信息，TTVSR采用Transformer的形式。
2. 对齐的本质就是将参考帧和支持帧进行内容上的匹配，理想情况就是两者一致。TTVSR推出的对齐方式和以往的不一样，他是基于patch的对齐，他会在每一个patch(即vision-token)的轨迹中寻找在时间域上最接近的(相当于最小化对齐误差)patch作为对齐结果，由于和过去或未来所有相关patch都相关，所以可以说它也是利用了时间信息；而且其可以利用更长时间上的信息，输入时间序列越长，则TTVSR对齐的搜索范围也越大，因此就可以获取最佳的对齐patch，故产生的对齐准确度也越高。TTVSR基于轨迹的对齐方法会让对齐只关注于同一个轨迹内的vision-token，而一个轨迹内的patch都是同一个物体在不同时间帧上的token，因此可以说TTVSR的对齐只关注和当前patch相关的特征，不仅如此，作者并不采用传统Transformer使用的加权平均获取新token的方式，而是使用了硬注意力机制，直接取最大权重的$V$(即value)来输出新的token。
3. 由于TTVSR对于注意力的计算只是在时间维度上，也就是Query中的每个patch只在同一个轨迹内的token之间计算自注意力，而并不像一些普通的基于Transformer的VSR(Vanilla-ViT)一样需要计算时间和空间上的注意力(如VSRT)，因此TTVSR在利用了全局相关性的同时大大降低了Transformer带来的计算复杂度问题。
4. 整个TTVSR实际用了3种对齐方式：①VSRT的特征传播(flow-based)；②预对齐(利用位图并行对齐多帧，也是flow-based)；③利用自注意力在轨迹上取相似度最高的vision-token来做patch-wise对齐。前两种可以看成是粗糙对齐，第三种进一步精确对齐。
5. 为了克服轨迹中同一物体分辨率变化的问题，作者推出cross-scale feature tokenization方法来保持空间不变性，不同于STN、DCN的做法去对抗空间变化，作者采用对不同尺度的池化窗口来保持空间不变性，具体为$4\times 4、6\times 6、8\times 8$，为了保持3种尺度下的最终输出图像大小一致，需要分别将图像先放大$1、1.5、2$倍。
6. TTVSR最终取得了SOTA的表现力，且模型量和运算复杂度都有个合适的值，具体如下：

Note：

1. 所谓的轨迹就是用位图(location map)记录同一个物体在不同帧上的patch位置信息，包括x轴和y轴的值。每一个patch都有1条轨迹，其由$t$个(每个batch的输入时间帧数目)$4\times 4$的patch组成，每一帧都有$(\frac{h}{4}\times \frac{w}{4})$个patch($h$、$w$分别为每个batch中每一帧的输入图像高和宽，实验中取$64\times 64$)。
2. 每次的相似度计算，$Q$只有在相同位置的patch帧中寻找，但由于这些帧都是预对齐过的，其本来的位置一般都会在别的地方(保存在位图中)，这些不同帧的patch连起来就是一条轨迹，轨迹展示的是同一个块内容的不断变化路径。
3. 和BasicVSR系列一样，为了解决特征信息利用的公平性问题，TTVSR也采用双向传播机制。
4. TTVSR的预对齐属于flow-based，但总体来看属于flow-free对齐。
5. 光流(flow)用$\Delta$表示(在高数或物理中，$\Delta$表示增量或变化量)，它本质就是向量场或者说就是运动矢量，代表相邻两帧之间的位移变化，因为光流估计就是运动估计，因此光流估计的结果就是网络预测的光流。此外，光流图的颜色代表运动方向，深浅代表运动速率，因此其就是运动矢量的表征。
6. patch-embedding：用来表示patch的向量，类似于NLP中的token-embedding，用来生成$K、Q、V$，其size为${\mathbb{R}}^{(c\ast 4\ast 4)\times (64/4\ast 64/4)}$。embedding-function是指代中途插入的映射函数，一般都是一些卷积层。
7. TTVSR的轨迹自注意力是用$\mathcal{Q}$中每一个token去查询$\mathcal{V}$中对应空间位置上在时间维度上最接近的的token，故torch.gather函数在做的时候选取时间之后对应空间位置的索取实际上取决于$\mathcal{Q}$。
8. 硬注意力矩阵的size是$b\times (64/4\ast 64/4)$，但实际用于torch.gather之前还需要复制$c\ast 4\ast 4$遍成$b\times 1\times (c\ast 4\ast 4)\times (64/4\ast 64/4)$，之所以可以直接复制是因为这256个token都是一个个整体，每一个整体内的最佳时间自然是需要一致的。
9. TTVSR的轨迹自注意力本质上属于one-stage对齐方式，其相当于一次光流估计和warp的过程。但是不同于后者，轨迹自注意力的对齐方式只关注目标token，故这种直接复制的方式可以减少后者带来的artifacts。

1 Introduction

在Transformer引入VSR之前，针对特征传播的不同，可以将VSR分为两个分支：①sliding-window；②Recurrent。
①Sliding-windows：参考帧的对齐只利用固定大小的时间尺度中的特征信息(即利用局部范围的信息)，比如5帧或7帧。典型代表有VESPCN、TDAN、Robust-LTD、EDVR等。
②Recurrent：参考帧的对齐会利用到全部输入视频的特征信息。典型代表有BasicVSR、BasicVSR++。

在BasicVSR中证明出利用全局的特征信息比只在固定时间窗口内利用局部特征信息更有利于视频重建，这是因为，过去或未来较远处的特征信息对于当前帧的对齐也是有帮助的，我们不能忽略；如下图所示是3种利用全局特征信息的VSR方法，他们都可以利用较远处的特征信息来恢复出图像的细节：
。如果sliding-window的VSR方法想要利用较远处的特征信息，就必须加大时间感受野(时间滤波器)的大小，但这样会直接增加计算复杂度。因此显然Recurrent这种特征传播方式更好，但是BasicVSR系列采用的双向循环网络虽然利用隐藏状态$\mathbf{h}$来传播特征从而让每一帧的对齐都利用到全局的特征信息，但RNN结构天然具有信息衰减、误差累积、以及缺乏利用长期信息的能力(梯度爆炸/消失)。

---

Transformer于2017年从NLP领域提出，因为其具有RNN不具备的并行处理特点，且其同样可以利用长距离的特征信息；不同于CNN使用卷积层的堆叠增加感受野以及RNN利用循环机制保存长序列信息，Transformer基于自注意力机制来捕捉长距离特征信息。

$TTVSR$
①Trajectory-Self-Attention
不同于VSRT计算时间和空间上的相关性，TTVSR只在经过预对齐的轨迹中计算相关性，也就是说其相关性的计算只在时间维度上，然后遍历所有时间，选取同一个位置上相似度最高的vision-token来作为对齐当前时间步上的patch；此外，每一个token都对应了1条轨迹，在实验中作者设置了$16\times 16=256$条轨迹。$Q$表示当前帧，其中每一个patch($4\times 4$)都去寻找最佳patch，对$Q$每一patch来说，它可以任选一个时间$t$下的feature map，这个时间$t$是根据$Q$和$K$计算所得最大相似度对应的时间所得；具体地，对于前向特征传播，$K$包含了过去所有时间的feature map，对于后向传播，$K$包含了未来所有时间的feature map。获取最佳时间之后，找到相应时间$t$下的$V[t]$，至于是这张feature map中哪一个patch，这是固定的，一定是和$Q$中那个patch同一个位置坐标下的patch，代码中用gather来控制。关于自注意力计算：①计算相似度的$K、V$已经是预对齐过的；②相似度运算和新token值的获取(torch.gather())都是在patch-wise上做的，通过将图像unfold成patch格式，自注意力运算结束再fold回图像格式${\mathbb{R}}^{b\times c\times 64\times 64}$；③自注意力计算没有设置softmax，是因为TTVSR只取最大权值的token，但为了避免相似度计算值过大就做了归一化处理(F.normalize())。

②Location Map
位图是为了产生预对齐轨迹用的。每一帧都有一张对应的位图，由于TTVSR将每张$64\times 64$的输入帧处理为$\frac{64}{4}\times \frac{64}{4}=16\times 16$的token，因此位图的size可表示为：${\mathbb{R}}^{b\times 2t\times 16\times 16}$，其中$2t$表示2个空间位置$x、y$轴和当前帧之前或之后的帧数。位图的更新和flow-based对齐是类似的(比如STN-Based)，唯一不同的是，这里copy的不是像素，而是坐标值！一份光流应用于多张位图；一份光流应用于多个通道；每个时间都配有1份位图(这里一份表示$x,y$两个通道)。预对齐的时候，$V$和位图的时间数必须要一样。
位图有以下性质：

1. 每张位图中保存的是该帧每个patch原始的$x、y$轴坐标。为什么是原始的坐标呢？因为位图会更新，每次对下一个新帧做对齐的时候，都要将当前时间步之前的所有位图全部更新；那么原始的坐标会被更新到哪里去呢？对于每张位图更新的结果就是将每个原始坐标保存到当前参考帧同一物体所在patch的位置上，具体示意图如下：。所谓预对齐就是对于在当前帧过去或未来所有帧中的每一帧，在它们的表达上($V$)以位图为STN的格点坐标，使用PyTorch的F.grid_sample()去$V$中过去或未来每一帧的表示上copy相应位置的patch来完成共$T-1$帧和当前参考帧的并行对齐，属于flow-based的patch-wise对齐，需要注意的是，经过卷积池化之后，物体的位置信息几乎不变(参考从RCNN到ATSS)。

2. 每张位图同一个位置上的坐标构成的就是该patch的一条轨迹，因为有$16\times 16$个位置，所以有$256$条轨迹。

3. 位图的更新是通过光流来做的(光流的计算是基于原始输入视频序列lrs，它和位图的size一样，是2维的张量${\mathbb{R}}^{b\times 64\times 64\times 2}$，这一点在VESPCN、Understanding Deformable in VSR(统一性的可变形卷积表达式、offset-fidelity损失函数)、VSRT都有说明)，但由于位图是基于patch的，所以还需要将光流flow进行池化成和位图一样的size：${\mathbb{R}}^{b\times 16\times 16\times 2}$。光流$flow\in {\mathbb{R}}^{b\times T-1\times 64\times 64\times 2}$的计算示意图如下所示(假设输入序列一共5帧，即$T=5$)：。在VSRT中，其多计算了1个时间维度的光流，但本质上还是和TTVSR是类似的：

4. 对于每一次当前帧的对齐，多张位图的更新都使用同一份光流(一般来说，对于通道数为$C$的张量，我们都只使用通道数为2的光流)。

5. 每张位图的初始化值均为其格点坐标。

③Cross-scale feature tokenization
由于一条轨迹中，同一个物体的大小会发生变化，最直观的就是同一个物体在某一帧中变大了，超出了patch的$4\times 4$范围，因此就需要增大感受野，然后使用池化手段来对抗物体的空间变换。通过使用全局平均池化来让多尺度下的图像统一成相同size的输出，即$V$有3种尺度下的表达形式，分别是使用$4\times 4、6\times 6、8\times 8$的窗口将图像分别放大$1、1.5、2$倍，具体是使用unfold-fold来扩大感受野。需要提一点的是，对于$4\times 4、8\times 8$在全局平均池化之后需要再使用一次和$4\times 4$一样的unfold-fold，最后都以图像格式保存而非patch格式，是因为接下去需要做flow-based的预对齐。

④Aggregation and SR-Reconstruction
在双向传播分支中，头和尾帧在刚开始直接和0特征图合并，相当于就只有头或尾帧；此外，双向分支并不是同时并行做的，而是先做完后向分支的对齐，然后做前向分支对齐，最后将2个分支的结果进行合并，然后通过SR重建网络输出超分的结果。
具体示意图如下(假设超分缩放倍率为4)：

---

Note：

1. 之所以用最大化$Q$和$K$的相似度计算结果作为索引来从$V$中取值，是因为$K$和$V$是同一个token的不同表达方式，可以理解为他们两是十分接近的，因此可以凭借$K$获取最佳时间步$t$来获取$V$在时间步$t$上相应位置的值。具体如下图所示：
   其中lr_curr_feat和feat_prop分别利用unfold-fold产生$K、V$，其中$V$有3种感受野的表达；$Q$也是由lr_curr_feat产生的，只不过其是下一个时间上的lr_curr_feat。

2. $Q、K、V$本质都是从$64\times 64$的图像中按$4\times 4$的patch分割(unfold)来的，具体如下所示(对于每个batch)：
   。

3. 无论是前向还是后向传播，假设$Q$为当前参考帧为从前向开始计算或从后向开始计算的第$T$帧($T\ge 1$)的表达，则$K、V$内包含了当前帧之前或之后所有帧的表达，即有$T-1$帧，$Q$总是领先$K、V$一个时间步。

4. 由于每个token的相似度计算都只发生在时间维度上，即$Q$中每个vision-token只和预对齐轨迹上，即不同时间上同一个空间位置处的patch计算相似度$\varphi$。这种方式避免了像VSRT那样同时在所有时间尺度和空间尺度上都进行相似度计算的高计算消耗。

5. 如何利用unfold-fold来增加感受野：增加感受野的方式，就是用较大kernel较小stride来产生较多密集的patch，然后用较大的kernel和较大的stride来fold。

6. TTVSR一共使用了3种对齐方式，除了自注意力和预对齐之外，其还使用了VSRT中相邻帧的flow-based对齐，这时候的支持帧是带有全局注意力的，因此和VSRT十分相似。这部分的对齐用于对自注意力对齐的进一步校正。

7. TTVSR有2处需要直接用到光流，分别是对齐校正和位图更新，它们都要使用mmediting库的flow_warp函数。但是mmediting的一个bug：flow_warp函数中应该是grid-flow而不应该是grid+flow。
   理由：①基于STN对于warp的理解；②VESPCN源码就是grid-flow；③反向采样原理，即表面的变换参数和实际warp的变换参数是相反的，光流估计由Spynet来做的，经过mmediting和一个类似Spynet源码的简易实现对比，这个没有问题，ta相当于表面的变换参数，而实际用于warp的变换参数因该对flow做相反变换，然后使用F.grid_sampling做warp才对。STN是在学习变换参数的时候，直接将变换参数往反向采样方向学习；④之所以+和-都对结果影响不大，是因为我们Spynet都是预训练的，比如先固定5000个回合，然后再预训练，他其实之后的回合都会去学习一个相反的光流，这就和我们直接使用相反光流是一样的。因次，两种方式只是开头不同，但那段影响不大。

8. TTVSR的光流估计使用SpyNet做的，反向光流的示意图如下：

9. 每个patch的size都是：${\mathbb{R}}^{(c\ast 4\ast 4)\times (h/4\ast w/4)}$，unfold虽然打乱了原图的空间位置，但是仍然保存原图中patch的空间位置信息的，且$Q、$