Snap Video: Scaled Spatiotemporal Transformers for Text-to-Video Synthesis

Paper name

Snap Video: Scaled Spatiotemporal Transformers for Text-to-Video Synthesis

Paper Reading Note

Paper URL: https://arxiv.org/pdf/2402.14797

Project URL: https://snap-research.github.io/snapvideo/index.html

TL;DR

• 2024 年 Snap Inc 提出的 Snap Video 方案。受 FIT 的启发提出了一种基于 transformer 的级联架构，相比于 Unet 有更好的可扩展性。同时针对视频生成任务优化了 EDM 中的 noise schedule 设计。根据研究参与者的反馈，Snap Video 在写实性方面与 Gen-2 相当，但显著优于 Pika。

---

Introduction

背景

• 目前文本生成视频的模型架构主要是在图像生成架构上插入特殊层来捕捉时间依赖性。同样，训练是在基于图像的扩散框架下进行的，模型被同时应用于视频和一组独立的图像，以提高结果的多样性。
• 我们认为这种方法在多个方面是次优的，这些方面在本工作中得到了系统地解决。
  • 首先，图像和视频模式在连续视频帧内容的相似性方面存在固有差异。类比而言，图像和视频压缩算法基于完全不同的方法。为了解决这个问题，我们重写了 EDM 框架，重点关注高分辨率视频。不像以往将视频处理为一系列图像的工作，我们通过将图像视为高帧率视频来进行联合视频-图像训练，以避免纯基于图像的训练中由于缺乏时间维度而引入的模式不匹配。
  • 其次，广泛采用的 U-Net 架构需要完全处理每个视频帧。与纯文本到图像模型相比，这增加了计算开销，并限制了模型的可扩展性（scaling）。scaling 是获得高质量结果的关键因素。扩展基于 U-Net 的架构以自然支持空间和时间维度需要体积 (volumetric) 注意力操作，这对计算的要求是极高的。无法做到这一点会影响输出，导致生成的动态图像或运动伪影，而不是具有连贯和多样化动作的视频。

本文方案

• 提出了 Snap Video 方案，利用帧间的重复性，并引入一种可扩展的 transformer 架构，将空间和时间维度视为单一的、压缩的 1D latent vector。这种高度压缩的表示使我们能够联合进行时空计算，并能建模复杂的运动。我们的架构受 FIT 的启发，并首次扩展到数十亿参数。与 U-Nets 相比，我们的模型在训练时间上减少了 3.31 倍，推理时间上减少了 4.49 倍，同时生成质量更高。
• 在广泛采用的 UCF101 和 MSR-VTT 数据集上评估了 Snap Video。我们的生成器在各种基准测试中表现出最先进的性能，特别是在生成运动质量方面。更有趣的是，我们对比了最近的开源和闭源方法，进行了多项用户研究，发现根据研究参与者的反馈，我们的模型在写实性方面与 Gen-2 相当，但显著优于 Pika 和 Floor33。最令人兴奋的是，在评估文本对齐和运动质量时，用户研究参与者大多数偏好 Snap Video。与 Gen-2 相比，我们的模型在提示-视频对齐方面在 81% 的情况下更受欢迎（对比Pika 为 80%，对比 Floor33 为81%），生成了最具动态的视频并具有最多的运动量（对比 Gen2 为96%，对比 Pika 为89%，对比Floor33 为88%），并且拥有最佳的运动质量（对比Gen-2 为79%，对比 Pika 为71%，对比 Floor33 为79%）。

Methods

EDM 介绍

EDM简介

扩散模型在图像和视频生成方面取得了显著的成功。EDM 定义了一个从方差爆炸 (variance-exploding) 的正向扩散过程到方差收缩的反向过程的去噪器函数。去噪器函数${\mathcal{D}}_{\theta }$通过去噪目标进行学习：

$$\mathcal{L}({\mathcal{D}}_{\theta })={\mathbb{E}}_{\sigma ,x_{\sigma },ϵ}\left[\lambda (\sigma ){\Vert {\mathcal{D}}_{\theta }(x_{\sigma })-x\Vert }_2^2\right],$$

其中$\lambda$是损失加权函数，$x\sim p_{\text{data}}$是一个数据样本，$ϵ$是高斯噪声，并且$\sigma$从训练分布中采样。${\mathcal{D}}_{\theta }(x_{\sigma })$被定义为：

$$\mathcal{D}({\tilde{x}}_{\sigma })={\mathcal{C}}_{\text{out}}(\sigma ){\mathcal{F}}_{\theta }({\mathcal{C}}_{\text{in}}(\sigma ){\tilde{x}}_{\sigma })+{\mathcal{C}}_{\text{skip}}(\sigma ){\tilde{x}}_{\sigma },$$

其中，${\mathcal{F}}_{\theta }$是一个神经网络，${\mathcal{C}}_{\text{out}}$、${\mathcal{C}}_{\text{skip}}$和${\mathcal{C}}_{\text{in}}$分别表示缩放函数。特别地，去噪目标$\mathcal{L}({\mathcal{F}}_{\theta })$可以等价地表示为：

$$\mathcal{L}({\mathcal{F}}_{\theta })={\mathbb{E}}_{\sigma ,x_{\sigma },ϵ}\left[w(\sigma ){\Vert {\mathcal{F}}_{\theta }({\mathcal{C}}_{\text{in}}(\sigma )x_{\sigma })-{\zeta }_{\text{nrm}}(\sigma ){\mathcal{F}}_{\text{tgt}}\Vert }_2^2\right],$$

其中，${\mathcal{F}}_{\text{tgt}}$表示训练目标，${\zeta }_{\text{nrm}}$是一个归一化因子，$w$是一个权重函数。这些形式在附录D中推导得出，并在表1中汇总。

提出了一个二阶龙格-库塔采样器来逆转扩散过程，并从高斯噪声中产生样本，其中$x_{{\sigma }_{\max }}\sim \mathcal{N}\left(0,{\sigma }_{\max }^{2}I\right)$。

改进 EDM 用于高分辨率视频生成

• EDM 最初被提出作为一种图像生成框架，其参数是为生成 64×64 像素的图像而优化的。空间分辨率的改变或引入帧间具有共享内容的视频，使去噪网络可以轻松地以原始分辨率恢复带有较高信噪比（SNR）的噪声帧，而原始框架设计为在较低噪声水平下工作。这是因为，考虑一个噪声视频 $x_{\sigma }\in {\mathbb{R}}^{T\times s\cdot H\times s\cdot W}\sim N(x,{\sigma }^{2}I)$，其中 T 是帧数，s 是上采样因子。我们通过对每个 T × s × s 像素块中的值取平均来构建相应的原始分辨率的干净帧和噪声帧 $\tilde{x},{\tilde{x}}_{\sigma }\in {\mathbb{R}}^{1\times H\times W}$。由于取平均的结果，噪声方差减少了一个因子 $T{s}^2$，即 ${\tilde{x}}_{\sigma }\sim N(\tilde{x},{\sigma }^{2}T{s}^{2}I)$，因此 ${\tilde{x}}_{\sigma }$ 的信噪比相对于 $x_{\sigma }$ 提高了（见图2）： $SN{R}_{{\tilde{x}}_{\sigma }}=T{s}^{2}SN{R}_{x_{\sigma }}$。如果每个块中的像素共享相似的内容（这是高分辨率视频中的典型情况），则平均帧中的信息对恢复 $x$ 是有用的，并且可以在训练时由去噪函数利用。这在初始采样步骤期间会产生训练和推理的不匹配，因为平均帧尚未包含一个形成良好的信号，而去噪器依赖其存在。因此，为了获得最佳性能，任何对 T 或 s 的更改应保持原始分辨率下相同的信噪比，这是扩散框架设计的目标。为了在原始分辨率下恢复最佳信噪比，可以通过一个对应的因子 ${\sigma }_{in}=s\sqrt{T}$ 来减少输入信号的幅度，如图2所示。因此，我们将前向过程重新定义为 $p(x_{\sigma }|x)\sim N(x/{\sigma }_{in},{\sigma }^{2}I)$。
  

• 重写 EDM，引入了输入缩放因子，并在表 1 中突出显示了这些变化。我们注意到，简单地引入缩放因子会以一种在小噪声值情况下使目标爆炸的方式改变训练目标 Ftgt。因此，我们利用以方程
  $$\mathcal{L}({\mathcal{F}}_{\theta })={\mathbb{E}}_{\sigma ,x,e}[w(\sigma ){\Vert {\mathcal{F}}_{\theta }\left(c_{\text{in }}(\sigma )x_{\sigma }\right)-c_{\text{mm }}(\sigma ){\mathcal{F}}_{\text{gt }}\Vert }_2^2$$
  形式表达的训练目标，以确保 Ftgt 保持不变的方式重写了 EDM 过程，使得有效损失权重 w(σ) 保持损失权重 λ(σ) 不变，cin(σ) 和 cnrm(σ) 将输入和训练目标规范化为单位方差，并且在 σin = 1 时框架等效于原始 EDM 表述。

最后，我们根据新定义的前向过程修改了采样器，该过程要求在 xσ 中的信号成分按 σin 进行缩放。通过将 Dθ(xσ) 除以 σin 并将最终去噪样本 x0 乘以 σin 来实现信号幅度的恢复。

图像-视频模态匹配

• 由于带字幕的视频数据相对于图像数据有限，联合图像-视频训练被广泛采用，并且通常对这两种模态应用相同的扩散过程。然而，如之前所述，视频中的 T 帧的存在要求与具有相同分辨率的图像采用不同的处理方式。一种可能的做法是对这两种模态采用不同的输入缩放因子。我们认为这种解决方案是不可取的，因为它增加了框架的复杂性，并且图像训练不会促进去噪模型学习时间推理，这是视频生成器的基本能力。为了在使用统一扩散过程的同时避开这些问题，我们通过将图像视为具有无限帧率的 T 帧视频来匹配图像和视频模态，并引入了一种可变帧率的训练过程，弥合图像和视频模态之间的差距。

模型架构，Scalable Video Generator

• U-Nets 在视频生成中取得了成功，通常通过增加时间注意力或卷积来建模时间维度。然而，这种方法需要对每个视频帧进行完整的 U-Net 前向传递，迅速变得成本高昂（见下图a）。这些因素对模型的可扩展性构成了实际限制，而可扩展性是实现高生成质量的主要因素，同样限制了时空联合建模的可能性。我们认为，将空间和时间建模分离会导致运动伪影、时间不一致或动态图像的生成，而不是生动运动的视频。然而，视频帧包含可压缩的空间和时间冗余内容。我们认为，学习压缩的视频表示，并联合建模空间和时间维度，是实现高质量视频生成所需的可扩展性和运动建模能力的必要步骤。
  

• FITs
  FITs 是最近提出的用于高分辨率图像合成和视频生成的高效 transformer 架构。其主要思想，如图 3 所示，是通过一组可学习的潜在 token 学习输入的压缩表示，并将计算集中在这一可学习的潜在空间上，使输入维度的增长对性能的影响较小。首先，FITs 对输入进行分块，并生成一系列块 token，然后将其分组。然后实例化一组潜在 token，并应用一系列计算块。每个块首先在潜在 token 和条件信号（如扩散时间步）之间执行交叉注意力“读取”操作，然后在相应组的潜在 token 和块 token 之间执行额外的组内“读取”交叉注意力操作以压缩块信息，应用一系列自注意力操作到潜在 token 上，并执行组内“写入”交叉注意力操作以将潜在 token 中的信息解压到块 token 中。最后，将块 token 投影回像素空间以形成输出。在一组潜在 token 上应用自调节以保留前几次采样步骤中计算的压缩视频表示。

• FITs 改进方案
  虽然有前景，但这些架构尚未扩展到最先进的基于 U-Net 的视频生成器的数十亿参数规模，也未应用于高分辨率视频生成。在下文中，我们强调实现这些目标所需的架构考虑。

  • 时间建模是高质量视频生成器的基本方面。FITs 通过考虑大小为 Tp×Hp×Wp​ 的三维块生成块 token，覆盖空间和时间维度。我们发现 Tp>1 的值限制了时间建模性能，因此我们只考虑覆盖空间维度的块。
  • 此外，与块 (patch) 类似，FITs 将块 token 分组，覆盖时间和空间维度，并逐组执行交叉注意力操作。我们观察到，为了实现最佳时间建模，每组的时间大小应配置为覆盖所有 T 视频帧。
  • 此外，由于时间维度的存在，视频相对于图像包含更多信息，因此我们增加了表示压缩空间的潜在 token 数量，以在其中进行时空联合计算。
  • 最后，FITs 使用局部层对相同组的块 token 执行自注意力操作。我们发现这种操作对大量块 token（我们的最大分辨率为 147456 个块 token）来说计算成本很高，因此在每次交叉注意力“读取”或“写入”操作后用前馈模块代替。

• 模型实现
  我们的模型使用由一系列条件 token 表示的条件信息来控制生成过程。除了表示当前 σ 的 token 外，为了启用文本条件，我们引入了一个 T5-11B 文本编码器，从输入文本中提取文本嵌入。为了支持训练数据中可变的视频帧率和大分辨率和纵横比差异，我们连接了额外的表示当前输入帧率和原始分辨率的令牌。为了生成高分辨率输出，我们实现了一个模型级联，包括生成 36×64 像素视频的第一阶段模型和生成 288×512 像素视频的第二阶段上采样模型。为了提高上采样质量，我们在训练期间用可变级别的噪声破坏第二阶段的低分辨率输入，并在推理期间通过超参数搜索对第一阶段的输出应用一定级别的噪声。

训练

我们使用 LAMB 优化器训练 Snap Video，学习率为 5e−3，采用余弦学习率 schedule，总批量大小为 2048 个视频和 2048 个图像，这得益于我们可扩展的视频生成器架构。我们在 55 万步上训练第一阶段模型，并在此基础上微调第二阶段模型，针对高分辨率视频进行微调，迭代 37 万次。我们设定 σin = s√T。考虑到 T = 16 帧的视频和为 EDM 设计的原始 64px 分辨率，我们将第一阶段的 σin 设为 4，将第二阶段模型的 σin 设为 32。

推理

我们使用 EDM 的确定性采样器和我们的两阶段级联从高斯噪声和用户提供的条件信息生成视频样本。我们在第一阶段使用 256 个采样步骤，在第二阶段使用 40 个采样步骤，除非另有说明，否则采用无分类器引导来改善文本视频对齐。我们发现动态阈值和振荡引导能持续提高样本质量。

Experiments

架构对比

• 相比于 U-Nets，FITs 效果更优，速度更快
  

噪声 scale 效果

• σin=4 的时候效果最好
  

与现有工作对比

• 取得 SOTA 效果
  

生成效果可视化对比

Thoughts

• 视频因为不同帧间可能存在冗余信息，原本为图像生成的噪声强度需要对应调整。调整思路是根据输入的分辨率和视频帧数计算一个噪声缩放系数，分辨率大的长视频需要加更多的噪声。
• 模型架构本文认为 spatial-temporal 架构不利于 scaling 的解释（将空间和时间建模分离会导致运动伪影、时间不一致或动态图像的生成，而不是生动运动的视频）有点奇怪，目前很多开源工作用 spatial-temporal 结构生成的视频运动也还挺好
• FITs 这种方式不知道对于分辨率、视频帧率变化过大的数据训练是否友好，需要模型自动处理不同压缩率下的学习