Paper name
W.A.L.T: Photorealistic Video Generation with Diffusion Models
Paper Reading Note
Paper URL: https://arxiv.org/pdf/2312.06662
Project URL: https://walt-video-diffusion.github.io/
TL;DR
- 2023 斯坦福大学和 google 联合出品的视频生成工作,提出了 Window Attention Latent Transformer (W.A.L.T),一种基于Transformer的潜在视频扩散模型(LVDMs)方法。训练了由三个模型组成的级联,包括一个基础的潜在视频扩散模型和两个视频超分辨率扩散模型,以每秒 8 帧的速度生成 512×896 分辨率的视频,并在 UCF-101 基准测试上报告了最先进的 FVD 得分。
Introduction
背景
- Transformers 在视频生成中的使用率一直不如 Unet,主要原因是全注意力的内存需求随着输入序列长度成平方增长。
- LDMs 这种潜在扩散模型从自动编码器的潜在空间中操作,降低了计算要求。需要选择是空间压缩还是时空压缩。
- 空间压缩可以用图像模型的权重初始化。但是增加网络复杂性和内存,限制了 transformer 的高分辨率视频生成
- 时空压缩可以缓解上述问题,但是不能用成对的图像-文本数据
本文方案
- 提出了 Window Attention Latent Transformer (W.A.L.T),一种基于Transformer的潜在视频扩散模型(LVDMs)方法。
- 自动编码器将视频和图像映射到一个统一的低维潜在空间,使得能够在图像和视频数据集上联合训练单一生成模型,并显著降低生成高分辨率视频的计算负担
- 提出了一种新的 Transformer 块设计用于潜在视频扩散建模,该块由交替的非重叠窗口限制的空间和时空注意力层组成
- 局部窗口注意力的使用显著降低了计算需求
- 促进了联合训练,空间层独立处理图像和视频帧,而时空层则专注于建模视频中的时间关系
- 训练了由三个模型组成的级联,包括一个基础的潜在视频扩散模型和两个视频超分辨率扩散模型,以每秒 8 帧的速度生成 512×896 分辨率的视频,并在 UCF-101 基准测试上报告了最先进的 FVD 得分
Methods
学习视觉 token
- 希望有一个共享且统一的压缩视觉表示,可以用于图像和视频的生成建模。第一帧总是独立于视频的其余部分进行编码。视频会在时序上下采样,类似 MAGVIT-v2。
学习生成视频和图片
h 和 w 维度上做 patchify 操作
- 加可学习的 position embedding,空间 position embedding 和时间 position emeddin 直接相加。图片的话就加视频第一个 latent 的 position embedding
窗口注意力:
- 空间窗口(SW)注意力仅限于潜在帧大小为 1×hp ×wp(第一维是时间)中的所有 token。SW 建模了图像和视频中的空间关系
- 时空窗口(STW)注意力仅限于大小为 (1+t) × h′p × h′w 的 3D 窗口内,建模视频潜在帧之间的时间关系。对于图像,我们仅使用 identity 注意力掩码,确保图像帧潜在值嵌入按原样通过层
- 除了绝对位置嵌入外,我们还使用相对位置嵌入(类似 swin transformer)。
条件生成
在 WALT 的 Transformer 骨干中,采用了三种类型的条件机制,如下所述:
- 交叉注意力。在我们的窗口 Transformer 块中的自注意力层之外,我们添加了用于文本条件生成的交叉注意力层。在仅在视频上训练模型时,交叉注意力层采用与自注意力层相同的窗口限制注意力,这意味着 S/ST 块将具有 SW/STW 交叉注意力层(图 2)。然而,在联合训练中,我们仅使用 SW 交叉注意力层。对于交叉注意力,我们将输入信号(查询)与条件信号(键、值)连接,因为我们早期的实验表明,这可以提高性能。
- AdaLN-LoRA。一种简单的方式来实现自适应层归一化是在每层 i 中包含一个 MLP 层,以回归条件参数向量 Ai = MLP(c + t),其中 Ai = concat(γ1, γ2, β1, β2, α1, α2),Ai ∈ R6×dmodel,c ∈ Rdmodel,t ∈ Rdmodel 是条件和时间步嵌入。在 Transformer 块中,γ 和 β 分别用于缩放和移位多头注意力和 MLP 层的输入,而 α 缩放多头注意力和 MLP 层的输出。这些额外的 MLP 层的参数数量随层数线性增长,并与模型的维度大小平方增长(块数 × dmodel × 6 × dmodel)。例如,在具有 10 亿参数的 ViT-g 模型中,MLP 层贡献了额外的 4.75 亿参数。受 LoRA 启发,我们提出了一种简单的解决方案,称为 AdaLN-LoRA,以减少模型参数。当 r ≪ dmodel 时,这显著减少了可训练的模型参数。例如,具有 r = 2 的 ViT-g 模型将 MLP 参数从 4.75 亿减少到 1200 万。
- 自我条件化。除了以外部输入为条件外,迭代生成算法还可以在推理过程中以自己先前生成的样本为条件。具体而言,Analog Bits 修改了扩散模型的训练过程,使得在某些概率 psc 下,模型首先生成一个样本,然后使用另一次前向传递,以此初始样本为条件进行精炼。以 1 − psc 的概率,仅进行一次前向传递。我们沿通道维度连接模型估计与输入,发现这种简单技术在结合 v-prediction 时效果良好。
自回归生成
- 为了通过自回归预测生成长视频,我们还将模型联合训练在帧预测任务上。这是通过在训练中以 pfp 概率对模型进行过去帧的条件化来实现的。具体而言,模型通过 cfp = concat(mfp ◦ zt, mfp) 进行条件化,其中 mfp 是二进制掩码。二进制掩码指示用于条件化的过去帧数量。我们在 1 个潜在帧(从图像生成视频)或 2 个潜在帧(视频预测)上进行条件化。这种条件化通过连接嘈杂潜在输入的通道维度集成到模型中。在推理期间,我们使用标准的无分类器指导,并将 cfp 作为条件信号。
视频超分辨率
使用单个模型生成高分辨率视频在计算上是不可行的。根据 cascaded diffusion models,我们使用级联方法,使用三个模型在不断增加的分辨率下进行操作。我们的基础模型生成分辨率为 128 × 128 的视频,随后通过两个超分辨率阶段进行两次上采样。我们首先使用 depth-to-space 的卷积操作在空间上扩展低分辨率视频。请注意,与训练中可以获得低分辨率的真实输入不同,推理过程中依赖于前一阶段生成的隐变量(参见教师强制)。为了减少这种差异,并提高超分辨率阶段在处理由较低分辨率阶段生成的伪影时的鲁棒性,我们使用噪声条件增强。具体来说,按照 γ(t) 添加噪声,通过从 U(0, tmax noise) 中采样噪声水平 tsr,并将其作为输入提供给我们的 AdaLN-LoRA 层。
长宽比微调
为了简化训练并利用具有不同长宽比的广泛数据源,我们在方形长宽比下训练我们的基础阶段。我们在一个数据子集上微调基础阶段,以通过插值位置嵌入生成 9:16 长宽比的视频。
Experiments
图像生成精度
视频生成精度
- 比 MAGVITv2 精度更高
消融实验
- patch size 越小精度越高
- window attention 和 full attention 精度差不多。window attention 能快 2x
- self condition 比例调高明显涨点
- adaln lora 影响,rank 越高参数量越大精度越好:
separate AdaLN 和 LoRA 对比:XL 模型 + LoRA 的 loss 更低,所以建议把更多参数量用在 base model 上,用 LORA 低 rank 就行
- zero snr 很重要
- autoencoder 的 latent dimension 越大,重建精度越高,但是生成精度不是 dimension 越大越好,生成的话 dimension 是 8 的时候最好。其实 16 看起来也不错
- image video 联合训练很重要,模型 scaling 涨点明显
生成效果
Thoughts
- 实验方面很充分,细节公开很多,非常多获取认知的 ablation 实验。很多模型设计都在计算量和效果之间平衡,比如用了两个超分模型和 adaln-lora。