SeedVR2:One-step video restoration via diffusion adversarial post-training

https://zhuanlan.zhihu.com/p/16547000551https://zhuanlan.zhihu.com/p/16547000551https://zhuanlan.zhihu.com/p/1917715064326162183https://zhuanlan.zhihu.com/p/1917715064326162183

高计算成本：基于扩散模型的视频修复在视觉质量上表现出色，但推理阶段计算开销巨大，限制了其实用性。

扩展困难：现有一步图像修复蒸馏方法难以直接扩展到视频修复任务，尤其是在处理高分辨率、真实场景视频时面临较大挑战。

窗口注意力的不一致问题：在高分辨率视频修复中，使用固定窗口大小的注意力机制容易导致窗口不一致，影响修复质量。

SeedVR2 模型：提出一种基于一步扩散的视频修复模型，通过与真实数据进行对抗训练，提升修复质量。

自适应窗口注意力机制：根据输出分辨率动态调整窗口大小，有效解决高分辨率下的窗口不一致问题。

对抗式后训练优化：在训练后期引入对抗训练框架，提升模型的逼真性与稳定性。

特征匹配损失设计：提出新的特征匹配损失，在不显著增加训练成本的前提下提升修复效果。

1.Introduction

基于SeedVR作为初始化，使用对抗训练目标对整个网络进行调优，消除了在蒸馏过程中从扩散教师预计算视频样本所带来的巨大成本。虽然直接采用 APT [34] 用于 VR 是适用的，但我们根据经验观察到，根据 VR 的性质，有几个关键方面可以改进。 首先 ，在以低质量输入为条件的情况下，我们观察到与文本到视频生成 相比，VR 的训练过程更稳定[34]， 即仅在单阶段的对抗训练中没有观察到明显的模式崩溃。但是，我们注意到在处理严重降级时性能下降。在此，我们在对抗训练之前采用渐进式蒸馏 [53]，以保持一步生成下的恢复能力。 其次，当在高分辨率 VR 上应用具有预定义窗口大小的窗口注意力时，例如，超过 2K 分辨率，我们观察到窗口补丁之间可见的边界伪影。我们推测这是由于窗口大小和训练分辨率的设置不当， 例如 ，与相对较小的训练分辨率相比，窗口大小过大，导致模型在处理窗口移位方面没有得到足够的训练。在处理不同分辨率的输入时，这种预定义的窗口方式可能会进一步限制每个窗口内 3D 旋转位置嵌入（RoPE） [63] 的鲁棒性。 为了解决这个问题，我们提出了一种自适应窗口注意力机制，在一定范围内动态调整窗口大小，从而显着提高了模型在处理任意分辨率输入时的鲁棒性。 第三 ，即使使用 APT，使用超大的生成器和鉴别器进行对抗性训练也可能不稳定， 即在长时间训练（ 例如 20k 迭代）后可以观察到性能下降。我们遵循 Huang 等 人。 [19] 通过引入 RpGAN [20] 和额外的近似 R2 正则化损失来增强训练稳定性。虽然 L1 损耗和 LPIPS 损耗 [92] 通常用于 VR 训练中，以更好地权衡 感知失真 [2]，但需要在像素空间中计算 LPIPS，这使得高分辨率视频训练无法负担。由于缺乏视频特定数据，训练潜在的 LPIPS 模型 [21] 也不适用。相反，我们提出了一种特征匹配损失来代替 LPIPS 损失，以实现高效的对抗训练。具体来说，我们直接从鉴别器的不同层中提取多个特征，并测量预测与地面实况之间的特征距离。我们根据经验表明，在我们的案例中，这种特征匹配损失是一种有效的替代方案。

2.Methodology

2.1 APT（Diffusion Adversarial Post-Training）

扩散对抗后训练是一种扩散加速方法，它将多步扩散模型转换为一步生成器，APT主要有两个阶段，即确定性蒸馏和对抗后训练。在确定性蒸馏中，首先在离散时间一致性蒸馏后训练蒸馏模型，并产生均方误差，教师模型生成蒸馏监督；对抗训练首先由预训练的扩散网络初始化鉴别器，然后引入额外的仅交叉注意力的transformer来生成用于损失计算的logits，为了稳定对抗训练，同时避免高阶梯度计算，APT提出了近似R1损失来正则化鉴别器，鉴别器的最终损失是非饱和GAN loss与近似R1损失相结合。

模型架构和自适应注意力窗口的分区。我们使用自适应窗口分区改进了 Swin-MMDIT [67]， 即通过调整大小的 LQ 输入（ ） Height×Width=960×960 上的 3×3 分区来确保窗口大小。用于计算特征匹配损失的特征是在 APT 中使用的交叉注意力层之前提取的[34]

2.2 Adaptive Window Attention

2.3 Training Procedures

在APT训练中并没有模式崩溃，但是在训练数千次迭代之后，可以观察到不必要的伪影，不稳定仍然存在，就目前测试来说，还是由不断闪烁的伪影。

Progressive Distillation，直接采用对抗训练从初始多步模型得到一步模型，由于初始模型与目标模型之间的巨大差异，可能会破坏模型的恢复能力，通过渐进式蒸馏来解决这个问题，具体来说，我们从从 SeedVR [67] 64 初始化的教师模型开始，采用采样步骤，然后逐步将学生模型提炼为一个步骤，蒸馏步幅为 2 。每个蒸馏过程需要大约 10 K 次迭代，并具有简单的均方误差损失。我们还在对抗训练期间逐步增加训练数据从图像到具有不同帧数的视频剪辑的时间长度，从而对各种长度的视频（包括图像）实现强大的 VR 性能。受益于这样的训练策略，我们进一步获得了从原始 7B 模型中提炼而成的 3B 模型，仅用模型大小的一半就实现了相当的性能。

Loss Improvement

3. Experiments

在72个H100-80G上训练SeedVR2，一个batch是100帧720p，采用序列并行和数据并行，每个阶段训练大概需要1天时间，首先按照本文中的新注意力设计从头训练一个7B的SeedVR模型，然后初始化7B SeedVR模型的参数，主要遵循APT中的训练设置进行训练，按照UAV合成了大约10M图像对和5M的视频对来训练。