MAGVIT: Masked Generative Video Transformer

Paper name

Paper Reading Note

Paper URL: https://arxiv.org/abs/2212.05199

Project URL: https://magvit.cs.cmu.edu/

Code URL: https://github.com/google-research/magvit

2023 年 CMU、google 等发表 CVPR2023 Highlight 文章，提出了视频生成方法 MAsked Generative VIdeo Transformer (MAGVIT)，基于两阶段方式训练，在多个视频生成测试集上取得了最佳效果。同时推理速度会显著优于同时期的 diffusion 方法和自回归方法。

受到 DALLE 等工作的启发。通过掩码 token 建模和多任务学习，提出了一种高效的视频生成模型
提出了 MAsked Generative VIdeo Transformer (MAGVIT)
- 第一个用于高效视频生成和操作的掩码多任务 transformer
- 单个训练模型可以在 10 种不同任务上推理
- 提出了一种高效的 embedding 方法，使用多样的掩码用于众多视频生成任务
- 在三个广泛使用的基准测试上取得了最佳的保真度性能，包括 UCF101，BAIR Robot Pushing 和 Kinetics-600
下图介绍
- a 展示了定量指标分析，实现了 SOTA 的 FVD 和 IS 指标，与之前最佳的 diffusion 模型（RaMViD, Video Diffusion）以及自回归模型（CCVS,TATS,NUWA）相比较
- b 展示了推理性能的优势，比 diffusion 模型快两个数量级，比自回归模型快 60x
  - 128 分辨率下，MAGVIT-B 在 V100 上可以达到 37fps，MAGVIT-L 在 TPU v4i 上可以达到 65fps
- c 展示了多个任务上的视频生成效果

MAGVIT 训练分为两个步骤
- 学习一个 3D 矢量量化（VQ）自编码器，将视频量化为离散标记
- 通过多任务掩码标记建模学习视频 transformer
下图展示了第二阶段的训练流程，在每个训练步骤中，随机选择一个任务及其提示 token，获取特定于任务的条件掩码，并优化 transformer 以在给定掩码输入的情况下预测所有目标 token：

基于非自回归 transformer，masked image synthesis 分为两个阶段
- 基于 Vector-quantized 自编码器将图像量化并展平为一系列离散序列
- 使用 masked token modeling （MTM）在离散序列上训练 transformer，训练目标是最小化掩码位置 token 与真实 token 之间的交叉熵
掩码图像生成在训练、测试阶段的具体流程（以 MaskGIT 为例）：
- 训练过程：基于一个余弦衰减的逻辑来确定掩码比例，来随机 mask 掉图片中的一些 token 进行训练
- 测试过程：使用非自回归的解码方式进行 12 步预测，从所有视觉 token 掩码的空白画布开始进行并行预测，每一步都并行预测所有 token，同时保留具有最高预测分数的 token，其余 token 被掩码并下一次迭代中进行预测，直到生成所有 token。

基于 VQGAN 改进
- VQ 自编码器是一个关键模块，它不仅为生成设置了质量界限，还确定了 token 序列长度，从而影响生成效率
- 现有方法在每帧上独立应用 VQ 编码器（2D-VQ）或在超体素上应用（3D-VQ），本文提出了一个不同的设计：将所有 2D 卷积扩展为带有时间轴的 3D 卷积。由于时间和空间维度的下采样率通常不同，使用 3D 和 2D 下采样层，其中 3D 下采样层出现在编码器较浅的层中，解码器在前几个块中使用 2D 上采样层，然后是 3D 上采样层
- 将 2D-VQ 的网络转换为带时间维度的 3D-VQ。同时使用 3D 膨胀 (3D inflation) 的方式，利用 2D-VQ 的权重初始化 3D-VQ。这对于 UCF-101 等小数据集较为有效。同时使用 reflect padding 替换 zeros padding，用于提高相同内容在不同位置的标记一致性
网络结构细节，其中灰色部分是主要不同的模块
训练细节
- 每帧使用 image perceptual 损失
- 基于以下优化使得 GAN loss 可以从头开始训
  - GAN loss 上增加了 LeCam regularization
  - 使用 StyleGAN 的 discriminator 架构，inflate 为 3D

采用各种掩码方案来进行训练，以适应具有不同条件的视频生成任务。这些条件可以是用于修复/生成图像的空间区域，也可以是用于帧预测/插值的几帧。
考虑十个多任务视频生成任务，其中每个任务具有不同的内部条件和掩码：帧预测（FP）、帧插值（FI）、中央外扩（OPC）、垂直外扩（OPV）、水平外扩（OPH）、动态外扩（OPD）、中央修复（IPC）和动态修复（IPD）、类别条件生成（CG）、类别条件帧预测（CFP）。
推理算法，固定推理步数进行非自回归预测
下图比较了非自回归图像解码（MTM， from MaskGIT）和本文的视频解码过程。与 MTM 解码不同，本文的解码从嵌入内部条件的多变量掩码开始，由此掩码引导，通过在每一步替换新生成的 token 的一部分进行有条件的转换过程，最终预测出所有 token，其中内部条件 token 得到了细化

在 MAGVIT 中评估了 3D-VQ 模型的不同设计方式。在 UCF101 数据集上对比。这里的指标衡量了中间量化的质量。结果表明：
- 尽管产生了更高的压缩率，3D-VQ模型显示出比2D-VQ更好的视频重建质量。
- 所提出的VQ在与相似大小的基线架构相比中表现得更好，并且在使用更大模型时效果更好
- imagenet 初始化涨点
- 中心膨胀优于平均膨胀