Open-Sora 1.1 技术报告

• 模型架构修改
• 支持不同视频长度/分辨率/宽高比/帧率（fps）训练
• 使用Masked DiT作为图生视频/视频生视频模型
• 数据收集和流程
• 训练详情
• 结果和评价
• 不足和下一步计划

在Open-Sora1.1版本中，我们使用了10M数据来训练经过结构调优后的STDiT的700M模型（Open-Sora1.0版本仅用400K数据）。我们实现了Sora报告中提到的以下功能：

• 可变的视频时长、分辨率、宽高比（包括采样灵活性、改进的取景范围和构图）
• 提示词增加图片和视频选项（使图像动起来、生成式增长视频、视频到视频编辑、连接不同视频）
• 图像生成功能

为了实现这一目标，我们在预训练阶段使用了多任务学习。对于扩散模型来说，用不同的采样时间步长进行训练已经是一种多任务学习。我们将这一思想在图像和视频的条件生成模型上，进一步扩展到多分辨率、宽高比、帧长、fps以及不同的掩码策略。我们在0~15s、144p到720p、各种宽高比的视频上训练模型。虽然由于训练FLOPs不足的限制，生成的视频在时间一致性上的表现没有那么高，但我们仍然可以看到这个模型的巨大潜力。

模型架构修改

我们对原始ST-DiT模型进行了以下修改，以获得更好的训练稳定性和模型性能（ST-DiT-2）：

• 在时间注意力模块中添加旋转位置编码：遵循目前LLM的最佳实践，我们将时间注意力模块中的正弦位置编码更改为旋转位置编码，因为它也算一项序列预测任务。
• 在时间注意力模块中添加AdaIN和Layernormal：我们将时间注意力与AdaIN和Layer范数作为空间注意力包裹起来，以稳定训练。
• QK归一化与RMSNorm：和SD3类似地，我们应用QK归一化来提高半精度训练的稳定性。
• 支持动态输入大小和视频条件限定：为了支持多分辨率、宽高比和fps训练，我们ST-DiT-2来接受任何输入大小。延申PixArt-alpha的想法，我们支持限定视频的高度、宽度、宽高比、帧长和fps。
• 将T5token数量从120扩展到200：我们使用的视频描述通常少于200个token，我们发现模型也可以很好地处理更长的文本。

支持不同视频长度/分辨率/宽高比/帧率（fps）训练

正如Sora报告中提到的，使用原始无损视频的分辨率、宽高比和视频长度进行训练可以增加采样灵活性，改善取景和构图。我们找到了三种实现这一目标的方法：

• NaViT：通过不同掩码策略支持在同一训练批次内使用不同大小的数据，并且训练效率下降很少。然而，该系统实现起来有点复杂，并且可能无法兼容kernal优化技术（如flashattention）。
• 填充（FiT，Open-Sora-Plan）：通过填充支持同一批次内的不同大小的数据。然而，将不同的分辨率填充到相同的大小会导致效率降低。
• 分桶训练（SDXL、PixArt）：支持通过分桶的方式在不同批次中动态调整大小，但在同一批次内数据大小必须相同，只能应用固定数量的数据大小。在一个批次中，我们不需要实现复杂的掩码或填充。

为了更便捷的实现，我们选择分桶训练的方式。我们预先定义了一些固定的分辨率，并将不同的样本分配到不同的桶中。下面列出了分桶方案中值得注意的点。但我们可以看到，这些在我们的实验中并不是一个大问题。

<details> <summary>查看注意事项</summary>

• 桶大小被限制为固定数量：首先，在实际应用中，通常只使用少数宽高比（9:16、3:4）和分辨率（240p、1080p）。其次，我们发现经过训练的模型可以很好地推广到未见过的解决方案。
• 每批的大小相同，打破了独立同分布（i.i.d.）假设：由于我们使用多个 GPU，因此不同 GPU 上的本地批次具有不同的大小。我们没有发现此问题导致性能显着下降。
• 可能没有足够的样本来填充每个桶，并且分布可能有偏差：首先，当本地批量大小不太大时，我们的数据集足够大以填充每个桶。其次，我们应该分析数据大小的分布并相应地定义桶大小。第三，分配不平衡并没有显着影响训练过程。
• 不同的分辨率和帧长可能有不同的处理速度：与PixArt只处理相似分辨率（相似token数）的宽高比不同，我们需要考虑不同分辨率和帧长的处理速度。我们可以使用“bucket_config”来定义每个桶的批量大小，以确保处理速度相似。

</details>

如图所示，桶是（分辨率，帧数量，宽高比）的三元组。我们为不同的分辨率提供预定义的宽高比，涵盖了大多数常见的视频宽高比。在每个epoch之前，我们打乱数据集并将样本分配到不同的桶中，如图所示。我们将样本放入最大分辨率和帧长度小于视频的桶中。

考虑到我们的计算资源有限，我们进一步为每个（分辨率，num_frame）二元组引入keep_prob和batch_size两个属性，以降低计算成本并实现多阶段训练。具体来说，高清视频将以概率1-keep_prob下采样到较低分辨率的桶中，并且每个桶的样本数量是由batch_size属性决定的。这样，我们可以控制不同桶中的样本数量，并通过为每个桶搜索合适的数据量来平衡GPU负载。

有关训练中桶使用的详细说明，请参阅配置文件.

使用Masked DiT作为图生视频/视频生视频模型

Transformer可以很容易地扩展到支持图生图和视频生视频的任务。我们提出了一种蒙版策略来支持图像和视频的调节。蒙版策略如下图所示。

在将图像或视频转换成另一个视频的过程中，我们通常会选择出需要作为条件的帧并取消其掩码（unmask）。在使用ST-DiT模型进行前向传播时，被选择取消掩码（unmask）的帧将被赋予时间步长0，而其他帧则保持它们原有的时间步长t。我们发现，如果直接将这种策略应用到训练好的模型上，会得到较差的结果，因为扩散模型在训练过程中并未学会如何处理一个样本中具有不同时间步长的帧。

受UL2的启发，我们在训练期间引入了随机掩码策略。具体来说，我们在训练期间随机取消掩码帧，包括取消掩码第一帧，前k帧，最后k帧，最后k帧，第一和最后k帧，随机帧等。基于Open-Sora 1.0模型，以50%的概率应用掩码策略，我们发现模型能够在10,000步的训练中学会处理图像条件（而30%的概率会导致处理能力变差），同时文本到视频的性能略有下降。因此，在Open-Sora 1.1版本中，我们从头开始预训练模型，并采用了掩码策略。

下图给出了用于推理的掩码策略配置的说明。五数字元组在定义掩码策略方面提供了极大的灵活性。

掩码策略用法的详细说明可在配置文件中查看.

数据收集和流程

正如我们在Sora1.0版本中看见的那样，数据数量和质量对于训练一个好的模型至关重要，因此，我们努力扩展数据集。首先，我们创建了一个遵循SVD的自动流水线，包括场景切割、字幕、各种评分和过滤以及数据集管理脚本和通用惯例。

我们计划使用panda-70M和其他数据来训练模型，大约包含3000万条数据。然而，我们发现磁盘输入输出（disk IO）在同时进行训练和数据处理时成为了一个瓶颈。因此，我们只能准备一个包含1000万条数据的数据集，并且没有完成我们构建的所有处理流程。最终，我们使用了包含970万视频和260万图像的数据集进行预训练，以及560,000视频和160万图像的数据集进行微调。预训练数据集的统计信息如下所示。

图像文本标记 (使用T5分词器)：

视频文本标记 (使用T5分词器)。我们直接使用Panda的短视频描述进行训练，并自己给其他数据集加视频描述。生成的字幕通常少于200个token。

视频时长：

训练详情

由于计算资源有限，我们必须仔细监控训练过程，并在推测模型学习不佳时更改训练策略，因为没有消融研究的计算。因此，Open-Sora1.1版本的训练包括多个更改，所以，指数移动平均（EMA）未被应用。

1. 首先，我们从Pixart-alpha-1024的模型checkpoint开始，使用不同分辨率的图像进行了6000步的微调。我们发现模型能够很容易地适应并生成不同分辨率的图像。为了加快扩散过程的训练，我们使用了SpeeDiT（iddpm-speed）技术。
2. [阶段一] 然后，我们使用梯度检查点（gradient-checkpointing）技术对模型进行了24,000步的预训练，这个过程在64个H800 GPU上运行了4天。尽管模型看到的数据样本数量相同，我们发现与使用较小批量大小相比，模型的学习速度较慢。我们推测，在训练的早期阶段，步数的数量对于训练更为重要。大多数视频的分辨率是240p，预训练时使用的配置与stage2.py相似。
3. [阶段一] 为了增加训练步数，我们改用了更小的批量大小，并且没有使用梯度检查点技术。在这个阶段，我们还引入了帧率（fps）条件。模型训练了40,000步，持续了2天。训练中使用的视频大多数是144p分辨率，使用的配置文件是stage1.py。我们使用较低的分辨率，因为我们在Open-Sora 1.0版本中发现模型可以以相对较低的分辨率学习时间知识。
4. [阶段一] 我们发现模型不能很好地学习长视频，并在Open-Sora1.0训练中发现了一个噪声生成结果，推测是半精度问题。因此，我们采用QK-归一化来稳定训练。我们还将iddpm-speed切换成iddpm。我们训练了17k步14小时。大多数视频的分辨率是144p，预训练时使用的配置是stage1.py。阶段1训练持续约一周，总步长81k。
5. [阶段二] 我们切换到更高的分辨率，其中大多数视频是240p和480p分辨率（stage2.py）。我们在所有预训练数据上训练了22000步，持续一天。
6. [阶段三] 我们切换到更高的分辨率，大多数视频的分辨率是480p和720p（stage3.py）。我们在高质量数据上训了4000步，用时一天。

结果和评价

不足和下一步计划

随着我们离Sora的复现又近了一步，我们发现当前模型存在许多不足，这些不足将在我们下阶段工作中得到改善。

• 噪音的生成和影响：我们发现生成的模型，特别是长视频中，有时很多噪点，不流畅。我们认为问题在于没有使用时间VAE。由于Pixart-Sigma发现适应新VAE很容易，我们计划在下一个版本中为模型开发时间VAE。
• 缺乏时间一致性：我们发现模型无法生成具有高时间一致性的视频，我们认为问题是由于缺乏训练FLOPs，我们计划收集更多数据并继续训练模型以提高时间一致性。
• 人像生成质量低：我们发现模型无法生成高质量的人类视频，我们认为问题是由于缺乏人类数据，我们计划收集更多的人类数据，并继续训练模型以提高人类生成。
• 美学得分低：我们发现模型的美学得分不高。问题在于缺少美学得分过滤，由于IO瓶颈没我们没有进行这一步骤。我们计划通过美学得分和微调模型来过滤数据，以提高美学得分。
• 长视频生成质量低：我们发现，使用同样的提示词，视频越长，质量越差。这意味着图像质量不能同等地被不同长度的序列所适应。

• 算法与加速实现：Zangwei Zheng, Xiangyu Peng, Shenggui Li, Hongxing Liu, Yukun Zhou
• 数据收集与处理：Xiangyu Peng, Zangwei Zheng, Chenhui Shen, Tom Young, Junjie Wang, Chenfeng Yu