视频生成Sora的从零复现：从Latte、Open-Sora(含1.0及其升级版)到StreamingT2V

v_JULY_v

已于 2024-04-27 23:53:12 修改

阅读量5.7k

点赞数 22

分类专栏： AIGC与CV多模态：图像生成、视频生成、自动驾驶文章标签：类Sora从零复现 Latte视频生成 Open-Sora 1.0

于 2024-03-19 15:51:36 首次发布

本文链接：https://blog.csdn.net/v_July_v/article/details/136845242

版权

AIGC与CV多模态：图像生成、视频生成、自动驾驶专栏收录该内容

10 篇文章 108 订阅

订阅专栏

前言

目前比较接近sora的开源路线是：Stable Video Diffusion(关于SVD的介绍请看此文的第4部分]) + Stable Diffusion3的结构(SD3的详细介绍见此文的第4部分)

其他的比如VDT虽然与sora架构最相似，但vdt本身因为没有做产品化所以效果比较一般(加之vdt的权重在24年Q1之内还没有开放)，如果想走这条路线，那需要做的改动还是比较大的
至于还有走基于videoGPT路线的，要做的就更多了，毕竟videoGPT都2021年的东西了

另，注意如果此前对sora原理并不了解的，请先看本sora系列的上一篇文章：视频生成模型Sora的全面解析：从AI绘画、ViT到ViViT、DiT、VDT、NaViT等(含sora复现)「下文提到此文时会简称为：sora全面解析一文」，避免一脸懵

第一部分 Sora爆火之前，上海人工智能实验室一团队推出Latte

24年1月5日，受DiT、U-ViT在图像生成领域的成功，上海AI实验室的一个团队推出了把“DiT/U-ViT架构用到视频生成领域”的类sora开源系统Latte(其对应论文为：Latte: Latent Diffusion Transformer for Video Generation，其对应的GitHub地址为：GitHub - Vchitect/Latte: Latte: Latent Diffusion Transformer for Video Generation.)

通过本博客之前的文章可知，stable diffusion所基于的潜在扩散模型(Latent Diffusion Models, LDMs)，通过在潜在空间而不是像素空间中进行扩散过程

LDMs首先利用预训练的变分自动编码器的编码器 $\mathcal{E}$ 将输入数据样本 $x \in p_{\text {data }}(x)$ 压缩成较低维度的潜在编码 $z=\mathcal{E}(x)$ 。随后，它通过两个关键过程学习数据分布：扩散和去噪

扩散过程逐渐将高斯噪声引入潜在编码 $z$ ，生成扰动样本 $z_{t}=\sqrt{\bar{\alpha}_{t}} z+\sqrt{1-\bar{\alpha}_{t}} \epsilon$ ，其中 $\epsilon \sim \mathcal{N}(0,1)$ ，遵循跨越 $T$ 个阶段的马尔可夫链
$\alpha_{t}$ 作为噪声调度器，其中 $t$ 表示扩散的时间步
去噪过程中，需要不断预测噪声 $z_{t-1}: p_{\theta}\left(z_{t-1} \mid z_{t}\right)=\mathcal{N}\left(\mu_{\theta}\left(z_{t}\right), \Sigma_{\theta}\left(z_{t}\right)\right)$ ，其中变分下界的对数似然缩减为 $\mathcal{L}_{\theta}=-\log p\left(z_{0} \mid z_{1}\right)+\sum_{t} D_{K L}\left(\left(q\left(z_{t-1} \mid z_{t}, z_{0}\right)|| p_{\theta}\left(z_{t-1} \mid z_{t}\right)\right) \mid\right.$
其中，噪声估计 $\mu_{\theta}$ 使用去噪模型 $\epsilon_{\theta}$ 实现，并通过以下目标函数进行训练(说白了，就是不断减少预测噪声与真实噪声之间的差距)
$\mathcal{L}_{\text {simple }}=\mathbb{E}_{\mathbf{z} \sim p(z), \epsilon \sim \mathcal{N}(0,1), t}\left[\left\|\epsilon-\epsilon_{\theta}\left(\mathbf{z}_{t}, t\right)\right\|_{2}^{2}\right]$

且为了使用学习的逆过程协方差Σθ来训练扩散模型，需要优化完整的 $D_{K L}$ 项，因此使用完整的 $L$ 进行训练，表示为 $\mathcal{L}_{v l b}$ 。此外， $\Sigma_{\theta}$ 是使用 $\epsilon_{\theta}$ 实现

更多细节请参见上文提过多次的此文第二部分《图像生成发展起源：从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer》

而Latte将LDMs扩展到视频生成，具体为：

首先通过预训练的变分自动编码器VAE $\mathcal{E}$ 将视频逐帧从像素空间压缩到潜在空间，然后视频 DiT 对隐式表征提取 token(提取视频patch组成tokens)
扩散过程在视频的潜在空间中进行，以建模潜在的空间和时间信息，其中噪声估计器 $\epsilon_{\theta}$ 通过transformer实现而非U-Net(即使用 Transformer 替换 U-Net 预测噪声的均值和方差，并且噪声是4维的)
过程中同时使用 $\mathcal{L}_{\text {simple }}$ 和 $\mathcal{L}_{v l b}$ 训练所有模型
最后 VAE 解码器将特征映射回像素空间生成视频

1.1 整体流程：输入、主干网络(Latte的4种变体)、输出

1.1.1 输入

视频片段经VAE压缩后假设为 $F \cdot H \cdot W \cdot C$ ，编码后得到 tokens $\hat{\boldsymbol{z}}$ ，维度为 $n_{f} \cdot n_{h} \cdot n_{w} \cdot d$
其中， $F, H, W,C$ 分别表示潜在空间中视频帧的数量、视频帧的高度、宽度和通道数(相当于总计F帧，然后每一帧的宽度、高度分别为H、W)
$\hat{\boldsymbol{z}}$ 和位置时空位置编码 $p$ 相加构成Transformer的输入 $z$

1.1.2 主干网络：Latte的4种变体(类似VDT)

主干网络：考虑到输入视频的空间和时间信息之间的固有差异以及从输入视频中提取出大量标记，如下图所示：从分解输入视频的空间和时间维度的角度，设计了4种高效的基于Transformer的模型变体(有没发现，下图的变体3类似于sora全面解析一文介绍过的VDT)，以有效捕获视频中的时空信息

首先是单注意力机制模块，每个模块中只包含时间或者空间注意力

时空交错式建模 (Variant 1)：先空间transformer、后时间transformer交替进行
如上图a所示，这种变体的Transformer骨干由两种不同类型的Transformer块组成：空间Transformer块和时间Transformer块
前者专注于在具有相同时间索引的token之间仅捕获空间信息，而后者以“交错融合”的方式跨时间维度以捕获时间信息

假设我们有一个视频剪辑在潜在空间中 $V_{L} \in \mathbb{R}^{F \times H \times W \times C}$ ，首先将 $V_{L}$ 转换为一个token序列： $\hat{\boldsymbol{z}} \in \mathbb{R}^{n_{f} \times n_{h} \times n_{w} \times d}$ (和Sora把视频内容转化成一系列token的思路完全一致)
视频剪辑在潜在空间中的token总数是 $n_{f} \times n_{h} \times n_{w}$ ，而 $d$ 表示每个token的维度，随后将时空位置嵌入 $p$ 加入到 $\hat{\boldsymbol{z}}$ 中，最后将得到的 $z=\hat{z}+p$ 作为Transformer骨干的输入

从而可以将 $z$ 由 $\hat{\boldsymbol{z}} \in \mathbb{R}^{n_{f} \times n_{h} \times n_{w} \times d}$ 重新整形成 $\boldsymbol{z}_{\boldsymbol{s}} \in \mathbb{R}^{n_{f} \times t \times d}$ (其中 $t=n_{h} \times n_{w}$ 表示每个时间索引的token数量)，作为空间Transformer块的输入，用于捕获空间信息
随后，包含空间信息的 $z_{s}$ 被重新调整成 $z_{t} \in \mathbb{R}^{t \times n_{f} \times d}$ ，作为时间Transformer块的输入，用于捕获时间信息

说白了，变体1的整个流程是首先将输入 reshape 成 $z_{s}$ ，维度为 $n_{f} \times t \times d$ ，其中 $t=n_{h} \times n_{w}$ ，先过空间Transoformer block，然后 reshape 成 $z_{t}$ ，维度为 $t \times n_{f} \times d$ ，过时间 Transformer block，循环几次
时空顺序式建模 (Variant 2)：先一组空间transformer 再一组时间transformer
与变体1中的时间“交错融合”设计相比，变体2采用“后期融合”方法来结合时空信息
如上图b所示，变体2有着与变体1相同数量的transformer块
与变体1类似，空间transformer块和时间transformer块的输入形状分别为 $\boldsymbol{z}_{\boldsymbol{s}} \in \mathbb{R}^{n_{f} \times t \times d}$ 和 $\boldsymbol{z}_{\boldsymbol{t}} \in \mathbb{R}^{t \times n_{f} \times d}$

说白了，变体2类似变体1，只是先一组空间再一组时序

接下来是，多注意力机制模块，每个模块中同时包含时间与空间注意力机制 (Open-sora所参考变体)

串联式时空注意力机制 (Variant 3)：只统一用一个transformer 先空间维度计算再时间维度计算
变体1和变体2主要关注Transformer块的因式分解，变体3专注于分解Transformer块中的多头注意力
如上图c所示，该变体首先在空间维度上计算自注意力，然后在时间维度上计算自注意力。因此，每个Transformer块都捕捉到了空间和时间信息
与变体1和变体2类似，空间多头自注意力、时间多头自注意力的输入分别为 $\boldsymbol{z}_{\boldsymbol{s}} \in \mathbb{R}^{n_{f} \times t \times d}$ 、 $\boldsymbol{z}_{\boldsymbol{t}} \in \mathbb{R}^{t \times n_{f} \times d}$

说白了，相当于a、b都使用完整的 Tranformer block 学习单一时间或空间信息，c是修改了Tranformer block 内部结构，MHA 先在空间维度计算自注意力，然后在时间维度计算
这样每个 Tranformer block 就相当于融合了时间和空间信息
并联式时空注意力机制 (Variant 4)：只统一用一个transformer 但拆成两个部分一部分计算空间一部分计算时间
在这个变体中，将多头注意力MHA分解为两个组件，每个组件利用一半的注意力头，如上图d所示
使用不同的组件分别处理空间和时间维度中的token，这些不同组件的输入形状分别为 $\boldsymbol{z}_{\boldsymbol{s}} \in \mathbb{R}^{n_{f} \times t \times d}$ 和 $\boldsymbol{z}_{\boldsymbol{t}} \in \mathbb{R}^{t \times n_{f} \times d}$
一旦计算出两个不同的注意力操作，我们将 $\boldsymbol{z}_{\boldsymbol{t}} \in \mathbb{R}^{t \times n_{f} \times d}$ 重新调整形状为 $\boldsymbol{z}_{\boldsymbol{t}}^{\prime} \in \mathbb{R}^{n_{f} \times t \times d}$ 。然后 $z^{\prime}$ 被添加到 $z_{s}$ 中，作为Transformer块中下一个模块的输入

说白了，相当于在 Tranformer block 内部将 MHA 分成了两部分，一部分计算空间，一部分计算时间，最后再融合

1.1.3 输出

输出：在Transformer主干之后，一个关键的过程涉及解码视频token序列，以得到预测的噪声和预测的协方差。这两个输出的形状与输入 $\boldsymbol{V}_{\boldsymbol{L}} \in \mathbb{R}^{F \times H \times W \times C}$ 相同

1.2 视频 patch embedding、Timestep-class 信息融入、空间位置 embedding、学习策略

1.2.1 视频 patch embedding

对于视频patch的嵌入，有如下两种方法

均匀帧patch嵌入
采样的 F帧都使用 ViT 的 patch 方式，类似平面操作，这样， $n_{f}=F, n_{w}=W / w, n_{h}=H / h$ ， $w$ 和 $h$ 表示每个 patch 的大小
压缩帧patch嵌入
使用 ViViT 的方式(ViViT的介绍见sora全面解析一文)，在时序上采样 $s$ ，类似立体操作(相当于捕捉时间信息，并将ViT块嵌入方法从2D扩展到3D，然后沿时间维度提取管道)，这样 $n_{f}=F / s$

很明显，第二种方法融入了时空信息，因此还需要使用3D转置卷积对“输出的潜在视频”进行时间维度的上采样，随后是标准线性解码器和重塑操作(an additional step entails integrating a 3D transposed convolution for temporal up sampling of the output latent videos, following the standard linear decoder and reshaping operation)

1.2.2 Timestep-class 信息融入

Timestep 以及 class 信息 $c$ 注入到模型采用了两种方式：

直接作为 tokens 加到输入中
使用 DiT 的自适应层归一化 AdaLN 方式
即基于输入 $c$ 通过 MLP 计算出 $\gamma_{c}$ 和 $\beta_{c}$ ，这样 $\operatorname{AdaLN}(h, c)=\gamma_{c} \text { LayerNorm }(h)+\beta_{c}$ ，其中 $h$ 表示为transformer快中的隐藏嵌入
此外，还新增一个 scale 向量 $\alpha_{c}$ ，对其进行回归操作，且应用在所有残差连接中，使得最终 $R C s(h, c)=\alpha_{c} h+\operatorname{AdaL} N(h, c)$ ，取名为 S-AdaLN

1.2.3 空间位置 embedding

时序位置编码方式一般有两种选择，比如绝对位置编码和相对位置编码RoPE(实验证明两种方式区别不大)，至于空间位置就是使用默认方式

1.2.4 学习策略：在DiT图像预训练模型基础上增加时间维度，且图像-视频联合训练

Latte 使用了DiT 在 ImageNet 上的预训练模型，并作了一些改进

DiT 位置编码 embedding 维度是 $n_{h} \cdot n_{w} \cdot d$ ，但是 Latte 的时空位置编码 embedding 维度是 $n_{f} \cdot n_{h} \cdot n_{w} \cdot d$ ，相当于因为需要带上时间维度的关系，在DiT的基础上复制了 $n_{f}$ 份
此外还去掉了 label embedding layer
且考虑到「基于 CNN 的视频生成方法提出图像和视频同时训练」会对最终的效果提升明显，故Latte将该技术应用到了Transformer 架构中
具体做法是在视频后面随机添加一些同一数据集其它视频的帧用于图像生成任务，时序相关的 tokens 只作用在视频部分

第二部分基于STDiT架构的类Sora模型Open-Sora 1.0及其升级版

2.1 Open-sora 1.0的架构设计

2.1.1 给文生图模型PixArt-α增加时间注意力层

模型采用了上面介绍过的 Diffusion Transformer(DiT)架构

作者团队以同样使用 DiT 架构的高质量开源文生图模型 PixArt-α [PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis] 为基座
然后在此基础上引入时间注意力层，将其扩展到了视频数据上

具体来说，整个架构包括一个预训练好的 VAE，一个文本编码器，和一个带有空间-时间注意力机制的STDiT模型(Spatial Temporal Diffusion Transformer，说白了，就是给DiT加上时空注意力，类似sora全面解析一文介绍过的VDT)，STDiT的初始化版本则相当于给PixArt-alpha 加上时间注意力

整个模型的训练和推理流程如下

在训练阶段，首先采用预训练好的 Variational Autoencoder (VAE) 的编码器将视频数据进行压缩，然后在压缩之后的潜在空间中与文本嵌入(text embedding)一起训练 STDiT 扩散模型
在推理阶段，从 VAE 的潜在空间中随机采样出一个高斯噪声，与提示词嵌入 (prompt embedding) 一起输入到 STDiT 中，得到去噪之后的特征，最后输入到 VAE 的解码器，解码得到视频

2.1.2 STDiT

再进一步阐述下STDiT，下图左侧为STDiT 每层的结构(看图时从下往上看)，下图右侧是上面2.4节介绍过的带Cross Attention的VDT，这两个结构的区别在于各个注意力层的组织顺序不同，但不影响本质是一样的

它采用串行的方式在二维的空间注意力模块上叠加一维的时间注意力模块，用于建模时序关系
在时间注意力模块之后，交叉注意力模块用于对齐文本的语意。与全注意力机制相比，这样的结构大大降低了训练和推理开销
与上节介绍过的同样使用「空间 - 时间注意力机制的 Latte (Latent Diffusion Transformer for Video Generation)」相比，STDiT 可以更好的利用已经预训练好的图像 DiT 的权重，从而在视频数据上继续训练

2.1.3 训练复现方案：类似SVD的三阶段训练

Open-Sora 的复现方案参考了Stable Video Diffusion的工作，共包括三个阶段，分别是：

基于「SD微调过后的图像VAE」做大规模图像预训练
第一阶段通过大规模图像预训练，借助成熟的文生图模型，有效降低视频预训练成本

具体而言，通过互联网上丰富的大规模图像数据和先进的文生图技术，可以训练一个高质量的文生图模型，该模型将作为下一阶段视频预训练的初始化权重。同时，由于目前没有高质量的时空 VAE(他们觉得VideoGPT的VQ-VAE质量较低，我也这么认为，所以在VideoGPT基础上去复现，会增加更多的工作 )，他们采用了Stability AI通过Stable Diffusion模型预训练好的图像 VAE [https://huggingface.co/stabilityai/sd-vae-ft-mse-original ]
基于「带时序注意力的文生图模型PixArt-alpha + T5」做大规模视频预训练
这个阶段需要使用大量视频数据训练，保证视频题材的多样性，从而增加模型的泛化能力(且有效掌握视频的时间序列关联)

Colossal-AI 团队使用了 PixArt-alpha 的开源权重作为第二阶段 STDiT 模型的初始化(且加载第一阶段权重，但与第一阶段不同的是加入了时序注意力模块，用于学习视频中的时序关系，同时初始化时序注意力模块输出为零，以达到更高效更快速的收敛)
以及采用了 T5 [T5: Text-To-Text Transfer Transformer] 模型作为文本编码器。同时他们采用了 256x256 的小分辨率进行预训练，进一步增加了收敛速度，降低训练成本
基于「高质量视频数据」进一步微调
第三阶段对高质量视频数据进行微调，显著提升视频生成的质量
作者团队提及第三阶段用到的视频数据规模比第二阶段要少一个量级，但是视频的时长、分辨率和质量都更高。通过这种方式进行微调，他们实现了视频生成从短到长、从低分辨率到高分辨率、从低保真度到高保真度的高效扩展

作者团队表示，在 Open-Sora 的复现流程中，他们使用了 64 块 H800 进行训练。第二阶段的训练量一共是 2808 GPU hours，约合 7000 美元，第三阶段的训练量是 1920 GPU hours，大约 4500 美元。经过初步估算，整个训练方案成功把 Open-Sora 复现流程控制在了 1 万美元左右

且每个阶段都会基于前一个阶段的权重继续训练。相比于从零开始单阶段训练，多阶段训练通过逐步扩展数据，更高效地达成高质量视频生成的目标

2.1.4 数据预处理

为了进一步降低 Sora 复现的门槛和复杂度，Colossal-AI 团队在代码仓库中还提供了便捷的视频数据预处理脚本，让大家可以轻松启动 Sora 复现预训练，包括公开视频数据集下载，长视频根据镜头连续性分割为短视频片段，使用开源大语言模型 LLaVA [https://github.com/haotian-liu/LLaVA] 生成精细的提示词

作者团队提到他们提供的批量视频标题生成代码可以用两卡 3 秒标注一个视频，并且质量接近于 GPT-4V。最终得到的视频 / 文本对可直接用于训练。借助他们在 GitHub 上提供的开源代码，我们可以轻松地在自己的数据集上快速生成训练所需的视频 / 文本对，显著降低了启动 Sora 复现项目的技术门槛和前期准备

2.2 Open-Sora的升级版

2.2.1 时空扩散模型ST-DiT-2

作者团队表示，他们对Open-Sora 1.0中的STDiT架构进行了关键性改进，比如

针对当前的序列预测任务，团队受LLM那边的启发，将时序注意力中的正弦波位置编码(sinusoidal positional encoding)替换为更加高效的旋转位置编码(RoPE embedding)
此外，为了增强训练的稳定性，他们参考SD3模型架构，进一步引入了QK归一化技术，以增强半精度训练的稳定性
为了支持多分辨率、不同长宽比和帧率的训练需求，作者团队提出的ST-DiT-2架构能够自动缩放位置编码，并处理不同大小尺寸的输入

2.2.2 多阶段训练

根据Open-Sora 技术报告指出，Open-Sora 采用了一种多阶段训练方法，每个阶段都会基于前一个阶段的权重继续训练。相较于单一阶段训练，这种多阶段训练通过分步骤引入数据，更高效地实现了高质量视频生成的目标

初始阶段大部分视频采用144p分辨率，同时与图片和 240p，480p 的视频进行混训，训练持续约1周，总步长81k
第二阶段将大部分视频数据分辨率提升至240p和480p，训练时长为1天，步长达到22k
第三阶段进一步增强至480p和720p，训练时长为1天，完成了4k步长的训练。整个多阶段训练流程在约9天内完成，与Open-Sora1.0相比，在多个维度提升了视频生成的质量

2.2.3 统一的图生视频/视频生视频框架

作者团队表示，基于Transformer的特性，可以轻松扩展 DiT 架构以支持图像到图像以及视频到视频的任务。他们提出了一种掩码策略来支持图像和视频的条件化处理。通过设置不同的掩码，可以支持各种生成任务，包括：图生视频，循环视频，视频延展，视频自回归生成，视频衔接，视频编辑，插帧等

作者团队表示，受到UL2[Ul2: Unifying language learning paradigms]方法的启发，他们在模型训练阶段引入了一种随机掩码策略

具体而言，在训练过程中以随机方式选择并取消掩码的帧，包括但不限于取消掩码第一帧、前k帧、后k帧、任意k帧等。作者还向我们透露，基于Open-Sora 1.0的实验，应用50%的概率应用掩码策略时，只需少量步数模型能够更好地学会处理图像条件化。在Open-Sora 最新版本中，他们采用了从头开始使用掩码策略进行预训练的方法

此外，作者团队还贴心地为推理阶段提供了掩码策略配置的详细指南，五个数字的元组形式在定义掩码策略时提供了极大的灵活性和控制力

2.2.4 支持多时间/分辨率/长宽比/帧率训练

OpenAI Sora的技术报告[3]指出，使用原始视频的分辨率、长宽比和长度进行训练可以增加采样灵活性，改善帧和构图。对此，作者团队提出了分桶的策略

具体怎么实现呢？通过深入阅读作者发布的技术报告，我们了解到，所谓的桶，是（分辨率，帧数，长宽比）的三元组。团队为不同分辨率的视频预定义了一系列宽高比，以覆盖大多数常见的视频宽高比类型。在每个训练周期epoch开始之前，他们会对数据集进行重新洗牌，并将样本根据其特征分配到相应的桶中。具体来说，他们会将每个样本放入一个分辨率和帧长度均小于或等于该视频特性的桶中。

作者团队进一步透露，为了降低计算资源的要求，他们为每个keep_prob和batch_size引入两个属性（分辨率，帧数），以减少计算成本并实现多阶段训练。这样，他们可以控制不同桶中的样本数量，并通过为每个桶搜索良好的批大小来平衡GPU负载。作者在技术报告中对此进行了详尽的阐述，感兴趣的小伙伴可以阅读作者在GitHub上发布的技术报告来获取更多的信息：https://github.com/hpcaitech/Open-Sora