[读论文] MaskINT: Video Editing via Interpolative Non-autoregressive Masked Transformers

Abstract

生成式 AI 的最新进展显著增强了图像和视频编辑，尤其是在文本提示控制的上下文中。最先进的方法主要依靠扩散模型来完成这些任务。然而，基于扩散的方法的计算需求很大，通常需要大规模的配对数据集进行训练，因此对实际应用中的部署提出了挑战。为了解决这些问题，本文将基于文本的视频编辑任务分为两个阶段。首先，我们利用预先训练的文本到图像扩散模型以零镜头的方式同时编辑几个关键帧。其次，我们引入了一个名为 MaskINT 的高效模型，它建立在非自回归 masked generative transformers 之上，专门使用中间帧的结构指导，专门研究编辑关键帧之间的帧细分。实验结果表明，我们的 MaskINT 与基于扩散的方法实现了相当的性能，同时显着缩短了推理时间。这项研究为基于文本的视频编辑提供了一种实用的解决方案，并展示了非自回归掩蔽生成转换器在该领域的潜力。

1. Introduction

paper: https://arxiv.org/abs/2312.12468

2. Related Work

3. Preliminaries

3.1 Masked Generative Transformers

Masked Generative Transformer 采用两阶段流程。第一阶段中，将图像通过向量量化（VQ）自动编码器量化为一系列离散的 token sequences。具体而言，给定图像 $I\in {\mathbb{R}}^{H\times W\times 3}$，编码器 $E$ 将其编码为一系列潜在向量，然后通过查找量化嵌入的 code-book 中的最近邻，将其离散化为有限个令牌。这样，图像可以表示为 code-book 的索引序列 Z = [ z i ] i = 1 h × w ,   z i ∈ { 1 , 2 , ⋯   , M } Z=[z_i]^{h\times w}_{i=1},~z_i\in\{1,2,\cdots,M\} Z=[zi​]i=1h×w​, zi​∈{1,2,⋯,M}，其中 $h$ 和 $w$ 为潜在特征的分辨率。解码器 $D$ 能够重构索引并近似原始图像，即 $D(Z)\approx I$

在第二阶段中，通过掩码令牌建模（Masked Token Modeling，MTM），训练一个双向 Transformer 模型。具体来说，在训练过程中，会随机选择一个掩码比例 $r\in (0,1)$，并将 $\lfloor \gamma (r)\cdot h\times w\rfloor$ 个 $Z$ 中的 token 替换为 [MASK] token，其中 $\gamma (r)$ 是 mask scheduling function。我们将带有掩码令牌的损坏序列表示为 $\bar{Z}$ ，并将类标签或文本提示等条件表示为 $c$。给定训练数据集 $D$，通过参数化的 BERT（参数为 $\Phi$） 来最小化每个掩码位置下的交叉熵损失：

$$L_{MTM}={\mathbb{E}}_{Z\in D}\sum _{{\bar{Z}}_i=[MASK]}-\log p_{\Phi }(z_i|\bar{Z},c)$$

在推理过程中，采用非自回归解码来生成图像。具体而言，图像或视频的所有令牌都初始化为 [MASK] 令牌。在第 $k$ 步，所有令牌并行预测，仅保留预测分数最高的令牌。剩余的分数较低的令牌将被再次掩码，并在下一次迭代中重新生成。每一步 $k$ 的掩码比例由 $\gamma (\frac{k}{K})$ 确定，其中 $K$ 为总迭代步数。

3.2 ControlNet

Latent Diffusion Models 是一类通过逐步去除噪声生成图像的去噪扩散概率模型。这种方法从一个初始的高斯噪声图像开始，经过 $T$ 个时间步逐步去除噪声，以生成最终图像。为了高效生成高分辨率图像，LDM 在自动编码器的潜在空间中进行扩散过程。具体而言，首先，编码器 $E^{\prime }$ 将图像 $I$ 压缩为低分辨率的潜在编码 $x=E^{\prime }(I)\in {\mathbb{R}}^{h\times w\times c}$。接着，带有注意力模块的 U-Net 模型 ${ϵ}_{\theta }$ 被训练用来去除噪声，其损失函数为：

$$L_{LDM}={\mathbb{E}}_{x_0,ϵ\sim \mathcal{N}(0,I),t\sim ||ϵ-{ϵ}_{\theta }(x_t,t,\tau )|{|}_2^2}$$

其中 $\tau$ 表示文本提示，$x_t$ 表示在 time step $t$ 的带噪声的潜在样本

ControlNet 针对仅使用文本提示难以准确描述生成图像的布局的问题，扩展了 LDM ，通过加入边缘图、深度图和人体姿态等空间布局条件来改善生成效果。具体来说，ControlNet 训练了与 LDM 相同的 U-Net 架构，并通过额外的条件对其进行微调，表示为 ${ϵ}_{\theta }^{\prime }(x_t,t,\tau ,s)$，其中 $s$ 表示空间布局。

4. Methodology

4.1 Overview

MaskINT 将视频编辑任务解耦为关键帧编辑阶段和结构感知的帧插值阶段。具体而言，给定一个包含 $N$ 帧的视频片段 { I i } i = 1 N \{I_i\}^N_{i=1} {Ii​}i=1N​，在第一阶段，利用输入的文本提示 $\tau$，同时对初始帧 $I_1$ 和最后一帧 $I_N$ 进行编辑，使用现有的图像编辑模型 $g(\cdot )$ 以零样本的方式进行。这一过程产生了高质量且连贯的编辑帧 $I_1^{\ast },I_N^{\ast }=g(I_1,I_N,\tau )$。

在第二阶段，MaskINT 通过非自回归变换器进行结构感知的帧插值。完整的编辑后视频帧由 { I i ∗ } i = 1 N = f Φ ( I 1 ∗ , I N ∗ , { S i } i = 1 N ) \{I^*_i\}^N_{i=1}=f_\Phi(I^*_1,I^*_N,\{S_i\}^N_{i=1}) {Ii∗​}i=1N​=fΦ​(I1∗​,IN∗​,{Si​}i=1N​) 生成，其中 $S_i\in [0,1{]}^{H\times W\times 1}$ 是结构条件。结构感知插值模块在仅包含视频的数据集上通过带掩码的标记建模（MTM）进行训练，利用结构信号 { S i } i = 1 N \{S_i\}^N_{i=1} {Si​}i=1N​ 以及初始帧 $I_1$ 和最后一帧 $I_N$ 进行条件生成。

4.2 Keyframes Joint Editing

根据文本提示 $\tau$ 及其边缘图 $S_1$ 和 $S_N$ 对关键帧 $I_1$ 和 $I_N$ 进行编辑。然而，即使对每个关键帧分别应用相同的噪声，ControlNet 仍无法保证外观一致性。为了解决这个问题，这篇文章扩展了自注意力机制以同时处理两个关键帧。具体来说，自注意力模块将时间步 $t$ 的第 $j$ 帧的噪声特征图 $x_j^t$ 投影到 $Q_j,K_j,V_j$ 上，然后在原始 U-Net 中进行处理。通过将所有选定关键帧的键和值拼接，扩展自注意力模块，使其在所有关键帧上执行注意力机制：

$$\text{Softmax}(\frac{Q_j[K_1,K_N{]}^T}{\sqrt{c}})[V_1,V_n]$$

由于每个帧都查询所有帧并从中汇聚信息，因此能够实现外观一致性。尽管公式中只使用了两个帧，但这种联合编辑可以扩展到任意数量的帧，但代价是处理时间增加和显著的 GPU 内存需求。

4.3 Structure-Aware Frame Interpolation

• Structural-aware embeddings. Non-autoregressive
  masked generative transformers 能够有效地执行帧预测和插值任务。然而，这些方法缺乏对结构的显式控制，使得跟踪输入视频的运动变得困难。因此，需要在生成过程中显式引入每一帧的结构条件。具体而言，使用现成的二维矢量量化（2D VQ）分词器对 RGB 帧 { I i } i = 1 N \{I_i\}^N_{i=1} {Ii​}i=1N​ 和结构图 { S i } i N = 1 \{S_i\}^N_i=1 {Si​}iN​=1 进行分词。选择 2D VQ 而非 3D VQ 以适应不同数量的帧和帧率。将 RGB 帧的分词表示为 c Z = { c z i } i = 1 N × h × w ^cZ=\{c_{z_i}\}^{N\times h\times w}_{i=1} cZ={czi​​}i=1N×h×w​（颜色分词），将边缘图的分词表示为 ${}^{s}Z={s_{z_i}}_{i=1}^{N\times h\times w}$（结构分词），其中 s z i , c z i ∈ { 1 , 2 , ⋯   , M } s_{z_i},c_{z_i}\in\{1,2,\cdots,M\} szi​​,czi​​∈{1,2,⋯,M}，$M$ 是码本大小。接下来，学习到的两个独立的嵌入层 $e_C(\cdot )$ 和 $e_S(\cdot )$ 分别将颜色分词和结构分词映射到各自的嵌入空间中。同时，添加了可学习的二维空间位置编码 $P_s$ 和和时间位置编码 $P_T$。因此，输入到后续层的内容可以表示为：

  $$X=e_c{(}^{c}Z)+e_s{(}^{s}Z)+P_S+P_T\in {\mathbb{R}}^{N\times h\times w\times c}$$

• Transformer with Window-Restricted Attention. 先前的 masked generative transformers 采用了具有全局注意力的 standard transformer。鉴于相邻帧之间没有显著的运动，我们在受限窗口内应用自注意力机制以进一步减轻计算开销。具体而言，该方法包括两个不同阶段的注意力。

  1. 首先，进行空间窗口注意力，在每一帧内部（尺寸为 $1\times h\times w$）计算注意力

  2. 然后，进行空间-时间窗口注意力，在一个尺寸为 $N\times h_w\times w_w$ 的管道内计算注意力，其中 $h_w$ 和 $w_w$ 是窗口大小

  为了进一步减少计算量，在编码层之前添加浅层卷积层对 $X$ 进行下采样，并在最后添加一个上采样层。

• Train. 使用仅包含视频的数据集来训练 MaskINT。令第 $i$ 帧的 RGB 帧颜色分词表示为 ${}^c{Z}_i$。在训练时，我们保留 ${}^c{Z}_1$ 和 ${}^c{Z}_N$ 的颜色分词，并随机将中间帧中 $[\gamma (r)\cdot (N-2)\cdot N]$ 个颜色分词替换为 [MASK] 标记。我们将这种损坏的视频颜色分词表示为 c Z ˉ = { c Z 1 , c Z ˉ 2 , ⋯   , c Z ˉ N − 1 , c Z N } ^c\bar Z=\{^cZ_1,^c\bar Z_2,\cdots,^c\bar Z_{N-1},^cZ_N\} cZˉ={cZ1​,cZˉ2​,⋯,cZˉN−1​,cZN​}。Structure-aware window restricted transformer（参数为 $\Theta$）的 loss 是：

  L M T M = E { c Z , s Z } ∈ D [ ∑ c z ˉ i = [MASK] − log ⁡ p Θ ( c z i ∣ c Z ˉ , s Z ) ] L_{MTM}=\mathbb E_{\{^cZ,^sZ\}\in\mathbb D}[\sum_{^c\bar z_i=\text{[MASK]}}-\log p_\Theta(^cz_i|^c\bar Z,^sZ)] LMTM​=E{cZ,sZ}∈D​[czˉi​=[MASK]∑​−logpΘ​(czi​∣cZˉ,sZ)]

• Inference. 在推理过程中，MaskINT 可以无缝地在联合编辑的关键帧之间进行帧插值，尽管它仅通过普通视频进行训练。具体来说，将第一阶段编辑后的初始帧 $I_1^{\ast }$ 和最后一帧 $I_N^{\ast }$ 分词为颜色分词 ${}^c{Z}_1^{\ast }$ 和 ${}^c{Z}_N^{\ast }$，并将所有中间帧的颜色分词初始化为 [MASK]。采用 MaskGiT 中的迭代解码，总共进行 $K$ 步。在第 $k$ 步，我们并行预测所有颜色分词，并保留置信度最高的分词。