083、VSRT：视频超分中的Transformer架构——时空注意力机制详解

083、VSRT视频超分中的Transformer架构——时空注意力机制详解上个月调一个视频超分模型跑了两天两夜loss曲线漂亮得像教科书上的示范图结果一上测试集相邻帧之间的闪烁能把人眼睛晃瞎。我盯着输出视频看了十分钟差点以为自己得了飞蚊症。后来排查到问题根源——模型压根没学会跨帧建模每一帧都在“自说自话”。这就是典型的时空注意力没做好的下场。今天聊VSRT这个把Transformer塞进视频超分的经典工作。不扯那些花里胡哨的公式推导咱们直接从代码和调参的视角把时空注意力这块硬骨头啃下来。为什么CNN在视频超分上总差一口气先说说为什么非得用Transformer。EDVR、BasicVSR这些CNN架构靠光流对齐可变形卷积在视频超分上已经打下一片江山。但有个致命伤——感受野受限。你叠再多的3D卷积相邻帧之间的长程依赖还是抓不住。比如一个物体被遮挡了两三帧又出现CNN基本就懵了要么产生鬼影要么直接丢失纹理。Transformer的self-attention天然适合干这事。它能在一帧内部捕捉全局依赖空间注意力也能跨帧建立联系时间注意力。VSRT的核心贡献就是把这两种注意力有机地拧在一起而不是简单粗暴地拼接。VSRT的时空注意力到底怎么设计的VSRT的架构可以拆成三块空间注意力模块、时间注意力模块、以及一个巧妙的交互机制。别被论文里的“交错式”这种词唬住说白了就是让空间和时间信息互相“串门”。空间注意力不是简单的ViT很多初学者直接拿ViT的patch embedding往视频帧上套结果发现显存直接爆炸。VSRT的做法更聪明——它把每一帧切分成非重叠的patch但注意这里有个关键设计不同帧的patch数量保持一致但位置编码是独立的。# 别这样写直接把所有帧拼成一个大序列# frames rearrange(frames, b t c h w - b (t h w) c) # 显存爆炸警告# 正确的做法逐帧处理保持帧间独立性defspatial_attention(self,x):# x shape: [B, T, C, H, W]B,T,C,H,Wx.shape# 每帧独立切patchpatchesx.unfold(3,self.patch_size,self.patch_size).unfold(4,self.patch_size,self.patch_size)patchespatches.contiguous().view(B,T,C,-1,self.patch_size,self.patch_size)# 这里踩过坑unfold之后一定要contiguous不然后面reshape会出玄学错误patchespatches.permute(0,1,3,2,4,5).contiguous()patchespatches.view(B,T,-1,C*self.patch_size*self.patch_size)# 对每帧独立做self-attentionfortinrange(T):patches[:,t]self.self_attn(patches[:,t])returnpatches空间注意力的核心是窗口限制。VSRT没有做全局self-attention而是用了局部窗口窗口大小一般是7x7或8x8。为什么因为视频帧的纹理细节通常是局部的全局注意力反而会引入噪声。这个设计后来被很多工作借鉴包括SwinIR。时间注意力跨帧对齐的优雅解法时间注意力是VSRT最出彩的地方。它没有用光流做显式对齐而是让模型自己学怎么“看”相邻帧。具体做法是对每个空间位置在时间维度上做self-attention。比如第t帧的某个patch它会去“关注”第t-1帧、第t-2帧对应位置的patch以及第t1帧、第t2帧的patch。这个范围时间窗口是个超参数论文里设的是5帧前后各2帧。deftemporal_attention(self,x,temporal_window5):# x shape: [B, T, N, D] N是patch数量D是特征维度B,T,N,Dx.shape half_windowtemporal_window//2# 这里有个坑边界帧的处理# 别这样写直接padding零# x_padded F.pad(x, (0,0,0,0, half_window, half_window)) # 会引入边界伪影# 正确的做法循环padding或者用相邻帧复制x_paddedtorch.cat([x[:,:half_window1].flip(1),x,x[:,-half_window-1:].flip(1)],dim1)output[]fortinrange(half_window,half_windowT):# 提取时间窗口内的特征windowx_padded[:,t-half_window:thalf_window1]# [B, temporal_window, N, D]# 在时间维度上做attentionattntorch.einsum(btn d, b t n d - b t n n,window,window)/math.sqrt(D)attnF.softmax(attn,dim-1)outtorch.einsum(b t n m, b m n d - b t n d,attn,window)output.append(out[:,half_window])# 只取中心帧的输出returntorch.stack(output,dim1)这个设计的精妙之处在于空间位置是绑定的。第t帧的patch只关注其他帧相同位置的patch这相当于做了一个隐式的光流对齐。模型通过注意力权重的大小自动学会哪些帧的对应位置信息有用哪些是遮挡或噪声。时空交互关键中的关键如果只是空间注意力时间注意力串行效果其实一般。VSRT的杀手锏是交叉注意力——让空间和时间信息在同一个特征空间里交互。具体实现上VSRT用了两个并行的注意力分支然后通过一个门控机制融合defspatio_temporal_fusion(self,spatial_feat,temporal_feat):# spatial_feat: [B, T, N, D] 空间注意力输出# temporal_feat: [B, T, N, D] 时间注意力输出# 门控机制让模型自己决定每个位置该信任哪个分支gatetorch.sigmoid(self.gate_conv(torch.cat([spatial_feat,temporal_feat],dim-1)))# 这里踩过坑gate的初始化很重要初始值偏0.5效果最好# 如果初始值偏0或1模型容易陷入局部最优fusedgate*spatial_feat(1-gate)*temporal_feat# 再加一层cross-attention让两个分支的信息充分混合fusedself.cross_attention(fused,fused,fused)returnfused这个门控机制解决了视频超分的一个经典问题静态区域和动态区域需要不同的处理策略。静态区域比如背景空间注意力就够了动态区域比如运动物体需要更多时间信息。门控让模型自己学会区分。训练中的那些坑VSRT的训练比普通超分模型要敏感得多。我踩过的坑包括学习率设置Transformer模块的学习率要比CNN模块小一个数量级。我试过统一用1e-4结果空间注意力模块直接发散。后来改成CNN部分1e-4Transformer部分1e-5才稳定下来。梯度裁剪这个一定要加。视频序列长梯度容易爆炸。我设的max_norm1.0效果不错。别设太小0.1的话模型基本学不动。时间窗口的选择不是越大越好。我试过7帧窗口显存直接翻倍效果反而下降了。原因很简单相隔太远的帧对应位置的内容可能已经完全变了强行让它们做attention只会引入噪声。5帧是个比较稳妥的选择。位置编码VSRT用了可学习的位置编码但要注意空间位置编码和时间位置编码要分开。我一开始图省事把时空位置编码合并成一个结果模型分不清“第1帧左上角”和“第5帧左上角”的区别。实战效果与局限性VSRT在REDS4和Vid4上的PSNR比EDVR高了0.3-0.5dB这个提升在视频超分领域已经相当可观。更关键的是它的输出视频几乎看不到闪烁时间一致性比CNN方法好太多。但VSRT也有明显的短板计算量太大。一个VSRT-Large模型处理30帧的720p视频需要24GB显存。这在实际应用中基本不可行。后来的一些工作如BasicVSR通过光流引导的注意力把计算量降了下来。对运动剧烈的场景处理不好。当物体运动速度超过每帧20个像素时间注意力的对齐效果就会变差。因为对应位置的patch可能已经完全不相关了。这个问题直到后来用光流做显式对齐才得到缓解。长视频的累积误差。VSRT是逐段处理的比如每5帧一组组与组之间没有信息传递。处理长视频时不同组之间的亮度、颜色可能会有跳变。这个问题的解决方案是加一个时序平滑损失或者用循环结构。个人经验总结如果你现在要做一个视频超分项目我的建议是别一上来就上VSRT。先跑通一个简单的CNN baseline比如EDVR确保数据流程、训练策略都没问题。VSRT的训练太敏感debug成本极高。如果非要用Transformer先从轻量级开始。VSRT-Tiny只有4个Transformer block参数量不到10M效果已经比大多数CNN方法好。没必要一开始就上Large版本。注意数据增强。视频超分的数据增强比图像超分更讲究。随机裁剪、水平翻转这些常规操作没问题但别用颜色抖动因为视频帧之间的颜色一致性很重要。我试过加颜色抖动结果模型学会了忽略颜色信息输出全是灰蒙蒙的。监控时间一致性。除了PSNR和SSIM一定要加一个时间一致性指标比如temporal flickering metric。很多模型PSNR很高但视频看起来就是不舒服就是因为时间一致性差。VSRT这篇工作虽然已经过去几年但它提出的时空注意力设计思路至今仍在影响视频超分领域。后来的很多工作比如Video Swin Transformer、TimeLens都能看到VSRT的影子。理解它的核心思想对你理解整个视频超分领域的发展脉络会有很大帮助。