发布 NVIDIA Cosmos 世界基础模型
摘要
我们推出 NVIDIA Cosmos™ 世界基础模型 (WFM)，这是一系列预训练模型，专门用于生成物理感知视频和世界状态，以推进物理 AI 开发。该系列包括用于文本到世界和视频到世界生成的自回归和扩散模型，为开发人员构建机器人、自动驾驶汽车和机器的世界模型提供了极好的基础。主要用例包括：
- 策略模型开发与评估
- 预测性预见建模
- 与 NVIDIA Omniverse 平台集成以进行多元宇宙模拟
为确保这些模型的安全使用，我们推出 Cosmos 护栏，这是一个最先进的系统，具有生成前和生成后防护功能，以维护提示完整性和输出一致性。
我们对 Cosmos WFM 的 3D 一致性和物理对齐进行了基准测试。Cosmos 模型在这些评估中始终优于基线视频合成模型。我们还公开提供 Cosmos 基准测试，以推进和评估未来的世界基础模型。
这些开放模型是 Cosmos 平台的一部分，该平台还包括数据整理工具、分词器和框架，可实现 Cosmos 世界基础模型的更快、更高效的微调。
快速链接
使用 Cosmos 世界基础模型 (WFM) 开发物理 AI
Cosmos 世界基础模型系列包括预训练的、专用模型，用于从文本、图像或视频输入生成物理感知视频和世界状态，从而推进物理 AI 开发。这些模型是能够捕获真实世界物理和自然行为通用知识的通用模型。
我们利用两种不同的可扩展深度学习范式：
- 扩散模型：将生成问题分解为一系列去噪任务。
- 自回归模型：将问题作为一系列下一个令牌预测任务来解决。
Cosmos 世界基础模型在以下条件下进行训练：
- 2000 万小时的视频数据（相当于 9000 万亿个令牌）
- 使用 10,000 个 NVIDIA H100 GPU，历时三个月
这些数据包括手部动作、物体操纵、空间感知、导航和相机运动。预训练整理确保模型针对推进机器人、自动驾驶汽车和其他物理 AI 系统进行了优化。
表 1：Cosmos 世界基础模型 1.0 版图示。我们有两套 WFM。一套基于扩散模型，另一套基于自回归模型。对于每个系列，我们构建了两个基础模型和两个派生模型。为了实现最佳生成质量，我们还为扩散模型构建了提示上采样器，为自回归模型构建了扩散解码器。
自回归模型
Cosmos 自回归模型利用输入文本、图像和过去的视频帧作为上下文，以更高的精度和速度预测未来的视频帧。
该架构专为物理 AI 用例量身定制，通过位置嵌入帮助增强生成控制，减少训练损失，并最大程度地减少视觉伪影。额外的交叉注意力层改善了文本理解，而归一化技术增加了稳定性——所有这些都确保了一致、更快且物理感知的真实输出。
图 1：。该流程首先通过 Cosmos-1.0-Tokenizer-DV8x16x16 的编码器对输入视频进行编码，生成离散令牌，这些令牌被转换为学习嵌入。这些嵌入通过重复的 Transformer 块进行处理，每个块都包含绝对位置嵌入和 3D RoPE 组件，这些组件在进入自注意力模块之前被展平。每个块还包括一个交叉注意力模块，该模块结合了编码的文本提示（通过 T5 文本编码器处理），然后是两层 MLP。最后，Cosmos-1.0-Tokenizer-DV8x16x16 的解码器从输出令牌重建视频。
Cosmos 自回归 WFM 的预训练
Cosmos 自回归世界基础模型的预训练遵循结构化的多阶段方法，以确保在视频预测和文本条件生成任务中具有强大的性能。
- 阶段 1：模型以视频预测目标开始，使用第一帧作为输入条件，训练预测未来 16 帧。此阶段使用 17 帧的上下文长度。
- 阶段 1.1：使用 YaRN 扩展增加上下文长度至 34 帧，以适应时间 RoPE，使模型能够捕获更长的视频依赖关系。
- 阶段 2：引入文本条件，通过新添加的交叉注意力层整合文本嵌入。模型在 34 帧上下文下进行训练，利用图像和视频数据进行联合训练。对于图像批次，由于其上下文长度较短，因此使用更大的批次大小。
- Cosmos-1.0-Autoregressive-4B：一个 4B 的 Transformer 模型，根据阶段 1 和阶段 1.1 的目标进行训练，用于下一个视频令牌预测。
- Cosmos-1.0-Autoregressive-5B-Video2World：派生自 4B 模型，它包含交叉注意力层，并进一步通过阶段 2 训练，用于文本条件视频生成。
- Cosmos-1.0-Autoregressive-12B：一个更大的 12B 的 Transformer 模型，根据阶段 1 和阶段 1.1 进行训练，专为高级的下一个视频令牌预测而设计。
- Cosmos-1.0-Autoregressive-13B-Video2World：基于 12B 模型构建，它包含交叉注意力层，并经过额外的阶段 2 训练，用于文本到视频任务。
- 减少训练和推理期间的计算成本
- 简化去噪任务
- Cosmos-1.0-Diffusion-7B-Text2World 和 Cosmos-1.0-Diffusion-14B-Text2World：根据文本描述生成 121 帧视频。
- Cosmos-1.0-Diffusion-7B-Video2World 和 Cosmos-1.0-Diffusion-14B-Video2World：根据文本描述和初始图像帧生成接下来的 120 帧。
- 预防护：使用黑名单检查和微调的 Aegis AI 内容安全模型扫描提示以查找不安全内容。
- 后防护：逐帧评估视频输出，拒绝不安全的视频。为了保护隐私和减少偏见，人脸会进行模糊处理。
- 通用模型：Cosmos 世界基础模型在真实世界物理和环境的各种数据集上进行训练，是通用模型，可以处理从自然动力学到机器人交互的广泛场景，为物理 AI 任务提供坚实的基础。
- 专业模型：开发人员可以使用更小、更有针对性的数据集微调通用模型，以创建适用于自动驾驶、人形机器人或自定义场景（如夜间应急车辆或工业机器人）的专业模型。与从头开始训练模型相比，微调可以减少数据和训练时间。
- 提速 89 倍的数据整理：显著减少处理时间。
- 可扩展性：处理超过 100 PB 的数据集。
- 高吞吐量：通过高级过滤、字幕和嵌入保持质量。
- 自回归模型：实现 8 倍时间压缩和 16x16 空间压缩，一次处理多达 49 帧。
- 扩散模型：实现 8 倍时间压缩和 8x8 空间压缩，处理多达 121 帧。
- 分片大型数据集以减少 I/O 开销。
- 确定性地保存和加载数据集，最大限度地减少重复和计算浪费。
- 优化通信以减少网络带宽使用。
- 3D 一致性
- 视觉一致性和保真度
- 物理对齐
模型以 640 × 1024 的固定分辨率进行训练。预训练后，进行冷却阶段，学习率在 30,000 次迭代中线性降低到零，同时在高质量图像-视频对上进行训练，从而精炼模型以实现高保真输出。
此版本中的模型变体
扩散模型
我们的基于扩散的 WFM 是潜在扩散模型，它们在分词器的学习潜在空间中操作，从而实现视频的紧凑、降维表示。这种设计选择具有以下几个优点：
为了将视频令牌化为潜在表示，我们采用 Cosmos-1.0-Tokenizer-CV8x8x8。
这些扩散模型的训练融合了多种先进技术，以优化性能和效率。3D 分块将视频或图像数据分解为不重叠的 3D 块，将其转换为网络的令牌序列，同时保留空间和时间关系。为了处理不同的视频大小、宽高比和帧速率，FPS 感知的 3D 旋转位置嵌入 (RoPE) 对跨时间、高度和宽度维度的位置信息进行编码，从而在渐进训练期间实现无缝适应。文本条件通过交叉注意力层实现，该层将来自 T5-XXL 嵌入的语义上下文与视觉令牌集成，以实现有效的文本到视频生成。查询-键归一化通过使用均方根归一化 (RMSNorm) 对注意力组件进行归一化来稳定训练，防止注意力崩溃等问题。此外，AdaLN-LoRA 通过自适应归一化中密集层的低秩近似将模型参数减少 36%（例如，从 11B 减少到 7B），在保持准确性的同时提高效率。总而言之，这些创新简化了训练，提高了视频生成质量，并实现了有效的基于文本的控制。
图 2：Cosmos-1.0-Diffusion 世界基础模型的整体架构。该模型通过 Cosmos-1.0-Tokenizer-CV8x8x8 的编码器处理输入视频以获得潜在表示，然后用高斯噪声对其进行扰动。然后，使用 3D 分块过程转换这些表示。在潜在空间中，该架构应用重复的自注意力、交叉注意力（整合输入文本）和前馈 MLP 层块，并通过自适应层归一化（比例、平移、门控）针对给定时间步长 𝑡 进行调制。Cosmos-1.0-Tokenizer-CV8x8x8 的解码器从精炼的潜在表示中重建最终视频输出。
Cosmos 扩散 WFM 的预训练
Cosmos 扩散 WFM 的训练方法旨在有效处理各种数据集、分辨率、宽高比和条件输入。联合图像和视频训练利用高质量图像数据集和视频数据，采用领域特定的归一化方案来对齐潜在分布并提高生成质量。渐进式训练从低分辨率视频（512p）开始，逐步过渡到更高分辨率（720p），并增加帧数，然后对高质量子集进行微调。多宽高比训练将数据组织到宽高比桶中（例如，1:1、16:9），并使用反射填充在调整大小时保留内容细节。混合精度训练通过保持 BF16 格式的权重以提高速度和 FP32 格式的权重以提高稳定性来优化效率，从而最大限度地减少损失峰值。文本条件将 T5-XXL 嵌入集成到 Text2World 模型中，确保提示和生成的视觉效果之间有很强的对齐。对于图像和视频条件，在训练期间将先前帧与生成的帧连接起来，并添加噪声以提高鲁棒性和灵活性。
此版本中的模型变体
使用 Cosmos 护栏确保安全
我们正在公开发布 Cosmos 护栏，以鼓励物理 AI 开发社区实现安全可靠的 AI。Cosmos 护栏分两个阶段运行：
图 3：Cosmos 护栏包括用于文本提示安全的预防护和确保安全视频输出的后防护。
如何使用 Cosmos 世界基础模型开发下游物理 AI 模型或自定义世界模型？
Cosmos 采用两阶段训练方法构建通用世界模型：
NVIDIA Cosmos 是一个世界基础模型开发平台，通过高效的视频处理、高性能分词器和先进框架简化训练，使开发人员能够快速有效地解决复杂的运营需求。
使用 NVIDIA NeMo Curator 加速数据处理
高质量数据对于训练模型至关重要，但准备起来可能很耗时。Cosmos 集成了针对 NVIDIA GPU 优化的 NVIDIA NeMo Curator，可高效处理海量数据集。例如，2000 万小时的视频可以在 NVIDIA Blackwell GPU 上仅用 14 天处理完成，而 CPU 管道需要 3.4 年。
主要优势
使用 Cosmos 分词器实现高保真压缩
Cosmos 分词器是一套用于图像和视频的视觉分词器，可在保持高重建质量的同时提供各种压缩率。Cosmos 分词器可以作为基于扩散模型和自回归模型中图像和视频生成的高效构建块。
这降低了成本和复杂性，同时保持了视觉质量，确保模型可以高效处理大型数据集。
使用 NVIDIA NeMo 框架进行微调
Cosmos 模型可以使用开放的 NVIDIA NeMo 框架进行微调，该框架可加速现有模型和新模型在 GPU 驱动系统上的训练。
这些工具使微调更快、更高效，无论是在本地还是在云端。
实时推理性能
Cosmos-1.0-Autoregressive-4B 模型在 8 块 NVIDIA H100 GPU 上提供了高效的推理性能。采用 320x512 分辨率和 10 FPS 视频设置，该模型处理 9 帧输入（0.9 秒，1280 令牌）以生成 24 帧未来帧（2.4 秒，1920 令牌），吞吐量为每秒 806 令牌，仅需 2.38 秒即可完成任务。主要开发工作包括：一个低分辨率分词器，在通用数据上预训练并针对 AV 和机器人等物理 AI 领域进行微调，以及使用此分词器微调自回归模型。此外，还使用 Medusa 优化了推测解码，并在 Alpamayo 数据集的 AV 数据上进行了微调。此设置确保了物理 AI 应用的高效率和高精度。
使用 Cosmos 基准评估世界基础模型
我们正在公开发布由 NVIDIA Research 与斯坦福大学和多伦多大学共同开发的 Cosmos 基准测试，以帮助物理 AI 社区评估世界基础模型。Cosmos 基准测试评估模型在机器人和自动驾驶汽车开发所需的 3D 一致性和物理对齐方面的表现。我们的第一代模型在这些指标上超越了基线 VideoLDM 世界模型。
使用 Sampson 误差评估几何精度，该误差测量兴趣点与其在另一个视图中的对应极线之间的距离的一阶近似值。较低的 Sampson 误差表示更好的几何理解。与基线模型 VideoLDM 相比，Cosmos 扩散模型和自回归模型实现了更低的 Sampson 误差和更高的姿态估计成功率，展示了卓越的几何理解能力。
使用峰值信噪比 (PSNR) 等指标进行评估，该指标测量图像中信号功率与噪声之间的比率。较高的 PSNR 值表示更好的视觉质量。Cosmos 模型始终优于基线，提供增强的时间一致性。
图 4：Cosmos 扩散模型和自回归模型都表现出较低的几何误差和较高的姿态估计成功率，突显了它们优于基线模型的几何理解能力。
虽然训练和评估物理对齐仍然是一项复杂的任务，但我们继续改进我们的模型，并使用通过像素、物体和特征级别指标评估直观物理的系统来测试它们是否符合物理定律。当以更多帧为条件时，Cosmos WFM 显示出改进的物体运动学预测。以提示 + 9 帧为条件的模型在像素级精度和视觉质量方面优于以提示 + 1 帧为条件的模型，这强调了数据整理和模型设计（可用于更好的物理对齐）的价值。
未来工作
我们通过持续训练和条件设置取得了改进的结果，始终超越基准并提高性能。展望未来，我们旨在将 NVIDIA Cosmos 与 NVIDIA Omniverse 等扩展物理 AI 平台集成，以解决机器人、自动驾驶汽车和机器中的实际挑战。
虽然 Cosmos 世界基础模型正在不断发展以实现更好的物理对齐，但我们认识到实现真正的现实世界物理理解的旅程仍在进行中。我们仍然致力于通过加强数据整理和改进模型设计来增强模型能力。
立即开始使用 NVIDIA Cosmos，并收听将于 1 月 7 日播出的由 NVIDIA 研究副总裁刘明宇主讲的 AI 播客。
