Sora - Open-Sora技术报告1.1
Sora - Open-Sora技术报告1.1
Open-Sora 1.1 是一个在视频生成领域取得重要进展的模型。相比其前身 Open-Sora 1.0，新版本在多个方面进行了重大改进，包括支持可变时长、分辨率和纵横比，以及更强大的图像和视频处理能力。本文详细介绍了 Open-Sora 1.1 的技术架构、训练过程和当前的局限性，为未来的研究和开发提供了有价值的参考。
在 Open-Sora 1.1 中，我们基于 10M 数据训练了一个参数量为 700M 的模型（Open-Sora 1.0 仅训练了 400K 数据），并采用了更先进的 STDiT 结构。我们实现了 Sora 报告中提到的以下特性：
- 可变时长、分辨率、纵横比（采样灵活性，改进的视频框架与构图）
- 图像和视频 prompt（图片动态化，视频扩展，视频编辑，视频连接）
- 图像生成能力
为了达成这些目标，我们在预训练阶段采用了多任务学习方法。对扩展模型，使用不同采样时间步已经是一种多任务学习。我们还将这个想法扩展到了多分辨率、纵横比、帧长度、fps，以及不同的图像视频条件生成掩码策略。我们使用 0~15s、144p 到 720p、不同纵横比的视频来训练模型。由于训练算力的限制，在时间一致性上的质量欠佳，但依然可以看到模型的潜力。
模型架构变更
我们将原始的 ST-DiT 进行了如下修改，称作 ST-DiT-2，以获得更好的训练稳定性和性能：
- Temporal Attention 中的 Rope Embedding：跟随 LLM 的最佳实践，将 Sinusoidal Positional Encoding 变更为 Rope Embedding。
- AdaIN 和 Layernorm 用于 Temporal Attention：用 AdaIN 和 Layernorm 包装 Temporal Attention，就像 Spatial Attention 一样，以稳定训练。
- 带有 RMSNorm 的 QK 归一化：跟随 SD3，将 QK 归一化应用于所有注意力，以提高半精度训练的稳定性。
动态输入大小支持和视频信息条件：为了支持多分辨率、综合比、fps 训练，SD-DiT-2 支持任意输入大小和自动伸缩位置编码。我们条件化了视频的高、宽、纵横比、帧长度、帧率，扩展了 PixArt-alpha 的理念。
将 T5 的 token 数从 120 增加至 200：文字说明通常少于 200 个 token，模型也较好地可以处理更长的文本。
支持多时间长度/分辨率/纵横比/帧率训练
就像 Sora 的报告里提到的，使用原始视频的分辨率、纵横比和长度进行训练可以增加采样灵活性并改进框架和构图。我们找到三种实现方法：
- NaViT: 使用掩码支持相同 batch 的动态大小，效率损失小。缺点是系统实现更复杂，不能获得 kernel 优化的收益（如 flash-attention）。
- Padding (FiT, Open-Sora-Plan)：使用填充支持相同 batch 的动态大小。缺点是将不同分辨率填充到相同大小效率低。
- Bucket (SDXL, PixArt)：使用分桶支持不同 batch 的动态大小。缺点是相同 batch 中的大小必须相同，且只能应用固定数量的大小种类。优点是同一 batch 中大小相同，不需要实现复杂的掩码或填充。
为了简化实现，我们使用了分桶方法。我们预定义了一批固定的分辨率，将不同的样本分到不同的桶中。分桶的主要担忧如下，可以看到这些问题在我们的用例场景下影响较小。
- 桶的大小是固定的种类：第一，在现实世界的应用中，常用的只有少数的纵横比(9:16, 3:4)和分辨率(240p, 1080p)；第二，训练后的模型可以很好地泛化到其他分辨率。
- 同一个 batch 内的样本大小固定，打通了独立同分布假设：我们使用了多个 GPU，GPU 的本地 batch 大小不同，我们没有观察到明显的性能下降。
- 可能没有足够的样本去填充所有的桶导致分布偏置：第一，我们的数据集足够大，在本地 batch 不是太大的情况下，足够填充所有的桶。第二，我们根据数据的大小分布来定义桶的大小。第三，不平衡的分布没有显著影响训练过程。
- 不同的分辨率和帧长度处理速度不同：不同于 PixArt 只能处理相似分辨率的纵横比（相同的 token 数量），我们需要考虑不同分辨率和帧长度的处理速度，可以通过 bucket_config 去定义每个桶的 batch 大小，保证相似的处理速度。
如图所示，桶是一个 (resolution, num_frame, aspect_ratio) 的三元组，我们提供了预定义常用的分辨率的纵横比。在每轮迭代前，打乱数据集并将样本分配到不同的桶中。如图所示。将一个样本放入比视频小的最大分辨率和帧长的桶中。
考虑到算力资源受限，对每对 (resolution, num_frame) 引入了 keep_prob 和 batch_size 参数来减小算力开销和使能多阶段训练。特别的，高分辨率的视频以 1-keep_prob 的概率被下采样到更低的分辨率。
每个桶的 batch 大小为 batch_size，通过控制不同桶的样本数，平衡了 GPU 之间的负载。
带掩码的 DiT 用作图像、视频生成模型
Transformer 可以轻易地扩展到图像到图像和视频到视频的任务。我们提出了一个掩码策略支持图像和视频的条件生成。掩码策略如下图所示：
典型的，我们在图像视频条件生成中，去除条件帧的掩码。在 ST-DiT 前向过程中，去掩码的帧的时间步为 0, 其他的会保持 t 不变。直接应用这个策略会导致较差的结果，由于扩散模型在训练中没有学会处理一个样本中的不同时间步。
受到 UL2 的启发，我们在训练中引入了随机掩码策略。特别地，我们在训练中随机对帧进行去掩码，包括首帧、前 K 帧，最后帧，最后 k 帧，最前和最后 k 帧，任意帧等。基于 50% 概率应用掩码的 Open-Sora 1.0，可以看到模型能够在 10k 步内学会处理图像条件（30% 的概率产生更差结果），伴随着生视频性能略降。因此，我们使用掩码策略重新预训练了 Open-Sora 1.1。
一个应用掩码策略配置的推理如下所示。一个五元组提供了定义掩码策略的灵活性。通过对生成的帧添加条件，我们可以自回归生成无限帧（尽管会传播误差）。
数据收集和流水线
我们在扩大数据集上做了很多工作，首先我们参考 SVD，创建了一个自动化的流水线，包括了场景切割，字母生成，各种评分和过滤，数据集管理脚本和规范。
我们计划使用 panda-70M 和其他数据训练模型，总计 30M+ 的数据。我们发现硬盘 IO 是数据处理和训练同时处理的瓶颈。因此我们只准备了 10M 的数据集且并没有经过我们所有的处理流水线。最后我们使用了 9.7M 的视频和 2.6M 的图像用于预训练，560K 的视频和 1.6M 的图像用于微调。预训练的数据集的统计数据如下：
训练细节
因为算力资源受限，我们仔细得观察了训练的过程，当预估模型学习效果不佳即没有进行消融实验的计算时，会调整训练策略。就是说，Open-Sora 1.1 的训练过程包含了很多改动，没有使用指数移动平均技术。
- 一开始，我们使用不同分辨率的图片，微调了 Pixart-alpha-1024 checkpoints 6k 步。发现模型可以轻易地生成不同分辨率的图片。使用了 SpeeDiT (iddpm-speed) 来加速扩散训练。
- [阶段1] 我们使用 gradient-checkpointing 技术预训练了模型 24K 步，在 64 个 H800 GPU 上花费了 4 天时间。尽管实际使用的样本数相同，但是相比更小的 batch size，模型学习得更慢。我们猜测在早期阶段，训练步数更重要。大部分的视频都是 240p 的分辨率。视频质量不错，但是模型并没有学习到很多时间知识，我们使用了 10% 的掩码率。
3.【阶段1】为了增加训练步数，我们切换到了更小的 batch size，并关闭了 gradient-checkpointing。我们加入了帧率条件控制，训练了 40k 步，花了 2 天时间。大部分的视频是 144p 的分辨率，我们使用更低的分辨率，因为我们在 Open-Sora 1.0 中发现，模型可以从低分辨率中学到时间知识。 - 【阶段1】视频不能很好地学习长视频，生成的结果存在噪声，我们推测原因和 Open-Sora 1.0 训练中发现的半精度问题有关。类似 SD3，模型可以快速适配 QK 归一化。同时将 iddpm-speed 切换至 iddpm，并将掩码比率提高到 25%，因为发现图像条件学习欠佳。我们训练了 17k 步，花费了 14 个小时。大部分的视频都是 144p 分辨率。至此第一阶段训练结束，整体花费了一周，训练了 81k 步。
- 【阶段2】转换到更高的分辨率，大部分视频是 240p 和 480p 的分辨率。训练了 22k 步，花费了 1 天。
- 【阶段3】转换到更高的分辨率，大部分视频是 480p 和 720p 的分辨率。训练了 4k 步，花费了 1 天。我们发现加载之前阶段的优化器状态，可以帮助模型学习得更快。
总结一下，Open-Sora 1.1 的训练在 64 个 H800 GPU 上，花费了 9 天时间。
限制和后续工作
离复现 Sora 更近了一步，同时也发现了很多当前模型的限制，这些限制指出了后续工作的方向：
- 生成失败：很多用例生成失败，特别是 token 总量比较大且内容比较复杂。在时间注意力机制中可能存在崩溃现象，且我们已经发现代码中一个潜在的 bug，我们正在努力修复它。此外，为了提升下一版本中生成内容的质量，我们将增大模型的规模并扩充训练数据。
- 含噪生成和不流畅：模型生成有时会含噪且不流畅，特别是长视频。我们认为问题在于没有使用 temporal VAE。鉴于 Pixart-Sigma 发现适配新的 VAE 很简单，我们计划在下一版本中为模型开发一个时间序列 VAE。
- 缺乏时间一致性：模型无法生成具有高时间一致性的视频。我们认为这个问题是由于训练过程中缺乏足够的浮点运算次数。我们计划收集更多的数据并继续训练模型，以提升时间一致性。
- 低质量的人像生成：模型无法生成高质量的人像视频。我们认为问题在于人像数据不足。我们计划收集更多的人像数据并继续训练模型，以提升人像生成的质量。
- 审美评分较低：模型的审美评分不高。问题在于缺乏审美评分的筛选，而这主要是因为 I/O 瓶颈。我们计划通过审美评分来筛选数据，并对模型进行微调，以提高其审美评分。
- 长视频生成质量更差：相同的 prompt，生成的视频越长，质量就越差。这意味着图像质量并未能同等适配不同长度的序列。