NeurIPS 2024 | VFIMamba:基于状态空间模型的视频插帧新SOTA

创业   2024-10-22 08:22   北京  

本文介绍一下我们(南京大学媒体计算组)与腾讯PCG合作的,并在最近被NeurIPS 2024接收的视频插帧方向的工作:VFIMamba: Video Frame Interpolation with State Space Models (VFIMamba)。
本文为了实现高效的全局动态时空建模,VFIMamba结合视频插帧任务的特点,将状态空间模型(SSMs)第一次应用于视频插帧任务中。为了适应状态空间模型需要对输入序列化的特征,我们提出了将首尾帧的token交错排列的方式来让SSMs更好的建模两帧的时空信息。

同时,为了更好地挖掘出SSMs模型在视频插帧任务的潜力,我们提出了基于课程学习思想的训练策略,通过逐渐增大训练集中的帧间运动幅度来增强SSMs对不同运动幅度下时空建模能力。在低分辨率插帧和高分辨率插帧的多个benchmark里,VFIMamba都取得了最好的性能。同时和之前的SOTA方法相比,VFIMamba在处理高分辨率插帧时的运行时间和占用内存都有明显的减少。

论文标题: 

VFIMamba: Video Frame Interpolation with State Space Models 

论文链接:

https://arxiv.org/abs/2407.02315

代码链接:

https://github.com/MCG-NJU/VFIMamba

一、引言

在视频插帧(VFI)中,帧间的时空建模至关重要,因为它被用于生成中间帧所需的运动估计和外观修复。目前的主流方法大多依赖于卷积或基于注意力的模型,这些模型往往要么缺乏足够的感受野,要么计算开销较大。最近,选择性状态空间模型(S6)应需而生,其特别针对长序列建模,提供了线性复杂度和数据依赖的建模能力。

在本文中,我们提出了VFIMamba,旨在通过利用S6模型实现高效且动态的帧间建模。我们的方法引入了混合状态空间模型块(MSB),该块首先以交错的方式重新排列相邻帧的标记,然后应用多方向的S6建模。这种设计不仅能有效传递帧间信息,还保持了线性复杂度。此外,我们还提出了一种新的课程学习策略,逐步培养在不同运动幅度下建模帧间动态的能力,充分发挥S6模型的潜力。实验结果表明,我们的方法在多项基准测试中达到了最先进的性能,尤其在高分辨率场景中表现突出。尤其是在X-TEST数据集上,VFIMamba在4K帧上实现了0.80 dB的显著提升,在2K帧上则达到了0.96 dB的改进

二、方法介绍

图1 模型结构图

2.1 整体结构

网络的整体结构如图1所示,我们首先通过一个轻量的特征提取器单独提取每一帧的信息。随后,利用本文提出的Mixed-SSM Block (MSB),将首尾帧拼接到一起并利用SSMs进行多方向的序列化的帧间时空建模。最后将得到的高质量的帧间特征用于中间帧的生成,其包含运动估计和外观修复两个部分。

2.2 如何将首尾帧融合为一个序列

因为SSMs仅限于处理一维序列,因此需要一种策略来扫描两个输入帧的特征图以进行帧间建模。在本文中,我们主要探索了不同的首尾帧融合方式,并给出了合理的解释。如图3,绿色表示首帧,红色表示尾帧。重新排列两个帧主要有两种方法:顺序排列(Sequential Rearrange),将帧串联成一个超图(Super Image);交错排列(Interleaved Rearrange),将两个帧的标记交错组合形成一个超图。

无论采用哪种重新排列方法,按照之前的研究,我们可以在得到的超图的四个方向上进行扫描:水平、垂直以及它们各自的反向方向。然后,SSMs被用于独立建模每个方向,最后将得到的序列重新排列并合并。
图2 不同的首位帧融合方式

通过分析,我们可以将SSMs的序列化建模中任意两个token i, j之间的建模表示为:

序列上不同token的相关系数
其中

表示序列中第 j 个标记 对第 i 个标记 的隐含注意力权重。与仅依据标记 (Qi) 和 (Kj) 的信息计算权重的注意力机制不同,SSMs 模型引入了 ,它包含了序列中第 i 个和第 j 个标记之间的上下文信息。基于这一结论,我们观察到在交错排列中,序列中任何一对标记的中间标记均来自它们的时空邻域。这意味着 结合了更多的局部建模,这对帧插值等低级任务是有利的。
此外,在交错排列中,时空相邻标记之间的中间标记数量通常较少。相比之下,在顺序排列中,即使是时空相邻的标记,在序列中也会被许多不相关的标记分隔。这可能引入噪声,干扰这些标记之间关系的建模。在图2 中可以看到一个具体的例子,第一帧的第 6 个标记和第二帧的第 11 个标记之间的标记在两种排列方法中差异显著。总之,我们认为对于视频帧插值,交错排列方法更适合进行更好的局部空间敏感处理。我们的实验进一步验证了这一结论。

2.3 基于课程学习思想的训练策略

尽管Mamba模型(S6)具有线性运算复杂度的全局感受野的优越特性,但通过适当的训练策略充分挖掘其潜力是更加至关重要的。目前,插帧算法主要采用两种训练策略:

1. Vimeo90K Only:大多数方法仅在 Vimeo90K 数据集上训练模型。尽管 Vimeo90K 提供了丰富多样的视频内容,但正如一些文章分析的那样,其数据集包含的运动幅度有限。这一限制妨碍了模型在大运动或高分辨率输入上的表现。

2. Sequential Learning:为了缓解仅在 Vimeo90K 上训练的局限性,一些方法在 Vimeo90K 初步训练后,进一步在 X-TRAIN 数据集上训练模型,该数据集以大运动和高分辨率内容为特征。尽管这种方法成功增强了模型在高分辨率数据上的性能,但往往导致对从 Vimeo90K 获取的小幅运动建模能力的遗忘。

为了解决这些问题并充分挖掘 S6 模型的潜力,我们提出了一种基于课程学习思想的训练策略,以学习不同运动幅度下的帧间建模能力,并同时保持对小幅运动的建模能力。具体而言,在继续对 Vimeo90K 进行训练的同时,我们逐步引入来自X-TRAIN 的数据。

X-TRAIN 的原始大小为 512×512,为了与 Vimeo90K 进行联合训练,我们首先将帧调整为 S×S 的大小,然后随机裁剪到与 Vimeo90K 相同的尺寸。每 T 个训练轮次,调整后的大小S 增加10%(从 256 开始),选定帧之间的时间间隔加倍(从 2 开始),这意味着随着训练的进行,运动幅度逐渐增加。这一策略使得模型能够逐步学习不同运动幅度下的帧间建模能力,从小幅运动开始,逐步过渡到较大幅度的运动。

三、实验结果

关于低分辨率的性能比较如图3:
图3 低分辨率数据集性能比较

关于高分辨率下插帧的性能比较如图4:

图4 高分辨率下性能比较

可以从结果看出来我们提出的方法在不同分辨率的数据集中都取得了SOTA的性能,其中在高分辨率的性能提升更加明显。图5是和之前SOTA方法的FLOPs和占用内存的比较,我们的方法随着输入尺寸的增大,计算开销有了成倍的减少:

图5 FLOPs和占用内存的比较

图6是视觉对比,我们提出的方法能够更好地生成中间帧:

图6 可视化对比
关于SSMs的消融:作为本研究的核心贡献,SSMs模型在帧间建模中实现了计算效率与高性能的平衡。为了验证其有效性,如图7 所示,我们进行了实验,移除了 MSB 中的 SSM 模型(w/o SSM),将 MSB 替换为 RIFE 中的卷/使用 EMA-VFI 中的局部帧间注意力/使用全局帧间注意力。我们观察到,仅移除 S6 模型导致参数减少了 0.7M,但在各个数据集上性能显著下降,凸显了 S6 的重要性。

在与卷积和局部注意力的比较中,我们发现虽然由于多个扫描方向,S6 模型相对较慢,但其性能提升显著。与全局注意力相比,S6 不仅超越了其性能,还提供了更快的推理速度和更低的内存消耗。总之,与现有模型相比,S6 模型确实在计算效率与性能之间实现了良好的平衡。

图7 SSMs有效性的消融
关于首尾帧融合方式的消融:输入帧如何排列对于使用 S6 模型进行帧间建模至关重要。正如在前面分析的那样,我们认为交错排列更适合 VFI 任务,并在此提供实验验证。如图8 所示,我们在水平和垂直扫描中实验了两种不同的重新排列方法。结果表明,使用交错排列在所有数据集上始终能取得最佳性能,并且相对于其他方法有显著提升。这些发现进一步验证了我们的分析,即交错排列在 VFI 中提供了更优越的时空局部建模能力。

图8 token排列方式的消融
关于训练策略消融:我们提出了一种课程学习的训练策略,以充分利用 S6 模型的全局建模能力。在图9 中,我们展示了不同学习策略在 Vimeo90K 和 X-TEST 数据集上随着训练轮次的性能表现。除了之前提到的 Vimeo90K Only 和 Sequential Learning 策略,我们还比较了一种基线方法,即将两个数据集直接混合进行训练(Mixed Learning)。
结果表明,随着训练轮次的增加,Vimeo90K Only 策略在 Vimeo90K 上的性能持续提升,而在 X-TEST 上变化甚微。Sequential Learning 最终虽然提高了 X-TEST 的性能,但在 Vimeo90K 上的表现显著下降。Mixed Learning 在两个数据集上的性能逐渐提高,但未能达到竞争水平。然而,我们提出的课程学习策略在训练结束时在两个数据集上同时实现了最佳性能。
图9 训练策略的消融

四、局限&未来展望

作为首个探讨在插帧任务中应用状态空间模型(SSMs)的研究,我们已取得非常高的性能,但仍然存在一些局限性。首先,尽管我们的方法比基于注意力的方法快得多,但仍未达到实时要求。因此,未来设计更高效的 SSM 将非常有价值。其次,在本研究中,我们主要关注 SSMs 在帧间建模中的作用,并未探讨其在帧生成模块中的应用。未来,直接使用 SSMs 生成中间帧也可能是一个值得探索的方向。
作者:张国珍
来源:【知乎】https://zhuanlan.zhihu.com/p/923110402
llustration From IconScout By IconScout Store

-The End-

扫码观看

本周上新!


“AI技术流”原创投稿计划


TechBeat是由将门创投建立的AI学习社区(www.techbeat.net社区上线500+期talk视频,3000+篇技术干货文章,方向覆盖CV/NLP/ML/Robotis等;每月定期举办顶会及其他线上交流活动,不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台,希望为AI人才打造更专业的服务和体验,加速并陪伴其成长。


投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //


投稿须知

稿件需要为原创文章,并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向,对用户启发更大的文章,做原创性内容奖励


投稿方式

发送邮件到

melodybai@thejiangmen.com

或添加工作人员微信(yellowsubbj投稿,沟通投稿详情;还可以关注“将门创投”公众号,后台回复“投稿”二字,获得投稿说明。


关于我“

将门是一家以专注于数智核心科技领域新型创投机构,也是北京市标杆型孵化器公司致力于通过连接技术与商业,发掘和培育具有全球影响力的科技创新企业,推动企业创新发展与产业升级。

将门成立于2015年底,创始团队由微软创投在中国的创始团队原班人马构建而成,曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业,不仅想获得投资,还希望获得一系列持续性、有价值的投后服务,欢迎发送或者推荐项目给我“门”: 
bp@thejiangmen.com

    


点击右上角,把文章分享到朋友圈
点击“阅读原文”按钮,查看社区原文

将门创投
将门是一家以专注于数智核心科技领域的新型创投机构,也是北京市标杆型孵化器,由前微软创投在中国的创始团队于2015年底创立。公司致力于通过连接技术与商业,发掘和培育具有全球影响力的科技创新企业,推动企业创新发展与产业升级。
 最新文章