25年1月来自 Nvidia 的技术报告“Cosmos World Foundation Model Platform for Physical AI”。
物理 AI 需要先进行数字化训练。它需要自身的数字孪生，即策略模型，以及世界的数字孪生，即世界模型。本文介绍 Cosmos 世界基础模型平台，以帮助开发人员为其物理 AI 设置构建定制的世界模型。本文将世界基础模型定位为通用世界模型，可以对其进行微调以适应下游应用的定制世界模型。该平台涵盖视频管理流程、预训练的世界基础模型、预训练的世界基础模型的后训练示例以及视频token化器。
如图所示 Cosmos World Foundation 模型： 预训练的 Cosmos WFM 可生成具有精确物理特性的高质量 3D 一致视频。Cosmos 模型套件包括扩散和自回归Transformer模型，它们分别使用视频的连续和离散潜在表示进行训练。使用专门的数据集对这些 WFM 进行后训练，使它们能够用于各种物理 AI 设置。具体来说，图中展示具有摄像头可控性的模型、能够遵循机器人操作指令的模型以及自动驾驶场景的模型。
物理AI是一种配备传感器和执行器的AI系统：传感器使它能够观察世界，执行器使它能够与世界互动并改变世界。它有望将人类工人从危险、费力或乏味的体力劳动中解放出来。尽管近十年来，由于数据和计算的扩展，AI的多个领域取得了长足进步，但物理AI却只取得了微小的进步。这主要是因为扩展物理AI的训练数据更具挑战性，因为所需的数据必须包含一系列交错的观察和操作。这些动作会扰乱物理世界，并可能对系统和世界造成严重损害。当AI仍处于起步阶段时，探索性行动至关重要，这种情况尤其如此。世界基础模型 (WFM) 是物理AI可以安全交互的物理世界的数字孪生，它一直是数据扩展问题长期以来寻求的补救措施。
本文介绍构建物理 AI 的 Cosmos 世界基础模型 (WFM) 平台。主要关注视觉世界基础模型，其中观察结果以视频形式呈现，并且扰动可以以各种形式存在。
如图所示：预训练的 WFM 是世界模型通才，它们使用大规模、多样化的视频数据集进行训练，捕捉现实世界物理的不同方面。这些预训练的世界基础模型可以通过后训练专门用于目标物理 AI 设置。通常，后训练的数据集是从目标物理 AI 设置收集的“提示”-视频对。提示可以是动作命令、轨迹、说明等形式。由于预训练的 WFM 提供良好的基础，因此后训练的数据集可以小得多。这种预训练和后训练为构建物理 AI 系统提供一种有效的策略。（图中虚线表示数据循环）
假设 𝑥/0:𝑡 为从时间 0 到 𝑡 对现实世界的一系列视觉观察。假设 𝑐/𝑡 为对世界的扰动。如图所示，WFM 是一个模型 𝒲，它基于过去的观察 𝑥/0:𝑡 和当前的扰动 𝑐/𝑡 预测时间 𝑡 + 1 的未来观察 𝑥ˆ/𝑡 +1。在例子中，𝑥/0:𝑡 是一个 RGB 视频，而 𝑐/𝑡 是一种可以采取多种形式的扰动。它可以是物理 AI 采取的动作、随机扰动、扰动的文本描述等。
如图所示 Cosmos WFM 平台中可用的内容，其中包括视频监管器、视频token化、世界基础模型预训练、世界基础模型后训练和护栏。
视频监管器
。开发一个可扩展的视频数据管理流水线。每个视频都被分成单独的镜头，没有场景变化。然后将一系列过滤步骤应用于剪辑，以定位高质量和动态信息丰富的子集进行训练。然后使用 VLM 注释这些高质量的镜头。然后执行语义重复数据删除以构建多样化但紧凑的数据集。
视频token化
。开发一系列具有不同压缩比的视频token化器。这些token化器是因果的。当前帧的token计算不是基于未来的观察。这种因果设计有几个好处。在训练方面，它使联合图像和视频训练成为可能，因为当输入是单个图像时，因果视频token化器也是图像token化器。这对于视频模型利用图像数据集进行训练非常重要，因为图像数据集包含丰富的世界外观信息，并且往往更加多样化。在应用方面，因果视频token化器与因果世界中的物理 AI 系统更加契合。
WFM 预训练
。探索两种可扩展的方法来构建预训练的世界基础模型——扩散模型和自回归模型。使用 Transformer 架构是因为它的可扩展性。
对于基于扩散的 WFM，预训练包括两个步骤：1)
Text2World
生成预训练，和 2)
Video2World
生成预训练。具体来说，训练模型根据输入的文本提示生成视频世界。然后，对其进行微调，以根据过去的视频和输入的文本提示生成未来的视频世界，将其称为
Video2World
生成任务。
对于基于自回归的 WFM，预训练包括两个步骤：1) 原始下一个 token 生成，和 2) 文本条件 Video2World 生成。首先训练模型，根据过去视频的输入生成未来的视频世界——预见（foresight）生成。然后，对其进行微调，根据过去的视频和文本提示生成未来的视频世界。
Video2world 生成模型是一个预训练的世界模型，它根据当前观察（过去的视频）和控制输入（提示）生成未来。对于基于扩散和基于自回归的 WFM，构建一系列具有不同容量的模型，并研究它们在各种下游应用中的有效性。
进一步微调预训练的扩散 WFM，以得到扩散解码器，从而增强自回归模型的生成结果。为了更好地控制 WFM，还基于大语言模型 (LLM) 构建一个提示上采样器。
世界模型后训练
。包括预训练的 WFM 在几个下游物理 AI 的应用。使用相机姿势作为输入提示对预训练的 WFM 进行微调。这能够在创建的世界中自由导航。针对人形机器人和自动驾驶任务对预训练的 WFM 进行微调。
护栏
。为了安全使用基础模型，开发一个护栏系统，可以阻挡有害的输入和输出。
给出一个视频监管流程，为token化器和 WFM 生成高质量的训练数据集。如图所示，该流程由 5 个主要步骤组成：1) 拆分、2) 过滤、3) 注释、4) 重复数据删除和 5) 分片。每个步骤都经过量身定制，以提高数据质量并满足模型训练的要求。
数据集
使用专有视频数据集和公开可用的开放互联网视频来训练模型。目标是为物理 AI 开发人员提供支持。为此，整理视频训练数据集以涵盖各种物理 AI 应用程序，并针对以下视频类别：
1. 驾驶 (11%)，
2. 手部动作和目标操控 (16%)，
3. 人体运动和活动 (10%)，
4. 空间觉察和导航 (16%)，
5. 第一人称视角 (8%)，
6. 自然动态 (20%)，
7. 动态相机运动 (8%)，
8. 合成渲染 (4%)，
9. 其他 (7%)。
拆分
视频长度不定，现代深度学习模型无法直接处理很长的视频。此外，许多视频包含镜头转换。它们可以从一个场景开始，然后转换到另一个场景，其中两个场景可以完全断开，例如，从两个人在纽约市现代厨房里交谈到狮子在非洲大草原追逐斑马的场景。根据每个视频的镜头变化对其进行分段并生成视觉上一致的视频剪辑非常重要，这样模型就可以学习物理上合理而非人工编辑的视觉内容转换。
拆分步骤生成的视频片段噪声很大，涵盖各种主题，质量差异很大。设计过滤步骤： 1) 删除视觉质量不符合最低要求的视频片段，2) 选择适合微调的高质量视频片段，3) 定制用于构建 WFM 的数据分布。通过运动过滤、视觉质量过滤、文本过滤和视频类型过滤来实现上述目标。
标注
文本描述通常与图像和视频数据配对，为世界模型训练提供监督和条件。使用 VLM 为每个视频片段生成高质量且一致的字幕。以某种方式配置 VLM，使其专注于视频中的重要事实和细节。使用这种方法提供视频描述而不是依赖 Alt 文本，也减轻世界模型的学习负担，因为不需要在训练期间适应不同的文本样式或格式。
测试几种 SOTA 方法，即
VFC
（Ge，2024）、
Qwen2-VL
（Wang，2024）、
VILA
（Lin，2024；Xue，2024）用于在视频上生成字幕，并发现 VILA 根据小规模的人工评估生成更准确的描述。使用具有 13B 参数的内部 VILA 模型，针对视频字幕进行微调。它具有扩大的上下文窗口，适合处理较长的多帧上下文，最大输入和输出 token 长度分别为 5904 和 256。为了提高推理效率，使用 FP8 量化的 TensorRT-LLM 引擎，与 PyTorch 半精度基线相比，吞吐量提高 10 倍，如表所示。提示 VILA“详细阐述视频的视觉和叙事元素”，并从输入剪辑中为其提供 8 个均匀采样的帧。字幕的平均长度为 559 个字符或 97 个单词。
去重
鉴于视频数量庞大，训练集中可能会有重复或近似重复的样本。对数据进行重复数据删除对于创建更平衡和多样化的数据分布至关重要。它还可以提高训练效率并减少记住特定训练样本的机会。
采用
SemDeDup
（Abbas，2023）和
DataComp
（Gadre，2024）的方法进行可扩展的语义重复数据删除。重用在过滤期间计算的 InternVideo2 嵌入，并使用多节点 GPU 加速的 k-means（RAPIDS，2023）实现对嵌入进行聚类，其中 𝑘 = 10,000。计算每个嵌入簇内的成对距离以识别重复项。当检测到重复的视频时，会选择分辨率最高的视频，以确保不会因重复数据删除而损失质量。为了避免将整个成对距离矩阵存储在 GPU 内存中，即时计算必要的上三角矩阵和 256 个块中的 argmax 约简。在重复数据删除过程中删除大约 30% 的训练数据。
还利用提取的 InternVideo2 嵌入和聚类结果构建一个可视化搜索引擎，支持使用自由格式的文本和视频查询整个训练数据集。该搜索引擎对于调试数据中的问题以及了解预训练数据集和下游应用程序之间的差距非常有用。
分片
此步骤旨在将处理后的视频片段打包成网络数据集，供模型训练器直接用于训练。根据视频的分辨率、宽高比和长度对视频进行分片，以符合训练课程。除了预训练数据集之外，还利用上面描述的不同过滤器创建质量更高的微调数据集。
基础设施
数据处理基础设施使用
AnyScale Ray
（Moritz，2017）为地理分布的集群实现一个流式管道系统，解决大规模 ML 工作流中的两个关键挑战：跨同质节点的高效资源利用以及对高延迟数据源连接的稳健操作。通过将数据传输与计算分离，管道可以高效地与远程数据存储配合使用，同时保持随管道复杂性而非数据集大小扩展的内存需求，从而实现无限制的流处理。
该架构通过并行管道阶段实现互补硬件资源的并发利用，例如，同时使用网络带宽进行数据提取、使用 NVDEC 单元进行视频解码以及使用 GPU 进行计算密集型转换。扩展碎片梯度下降算法（Weng et al.，2023）来优化这种多资源分配，该调度程序会自动扩展各个阶段，以在专用硬件加速器之间保持平衡的吞吐量。
Token化器是现代大模型的基本构建块。它们学习以无监督方式发现的瓶颈潜空间，将原始数据转换为更高效的表示。具体而言，视觉token化器将原始和冗余的视觉数据（例如图像和视频）映射到紧凑的语义token中，这使得处理高维视觉数据至关重要。这种能力不仅可以高效地训练大规模 Transformer 模型，还可以使它们在有限的计算资源上进行推理民主化。如图示意性地说明了token化训练流水线，其目标是训练编码器和解码器，以便瓶颈token表示最大限度地保留输入中的视觉信息。
Token化器有两种类型：连续和离散。连续token化器将视觉数据编码为连续潜嵌入，如SD（Rombach，2022）或 VideoLDM（Blattmann，2023）等潜扩散模型。这些嵌入适用于通过从连续分布中采样来生成数据的模型。离散token化器将视觉数据编码为离散潜代码，并将其映射到量化索引中，如
VideoPoet
(Kondratyuk et al., 2024) 等自回归Transformer中所示。这种离散表示对于使用交叉熵损失训练的模型（例如 GPT）是必需的。如图说明两种类型的token：
Token化器的成功很大程度上取决于它们在不影响后续视觉重建质量的情况下提供高压缩率的能力。一方面，高压缩率可以减少存储和计算需求。另一方面，过度压缩会导致重要视觉细节的丢失。这种权衡对token化器设计提出了重大挑战。
Cosmos Token化器，这是一套视觉token化器，其中包括用于图像和视频的连续和离散token化器。Cosmos Token化器提供卓越的视觉重建质量和推理效率。它提供一系列压缩率，以适应不同的计算约束和应用需求。
本文使用具有时间因果机制的轻量级、计算效率高的架构设计 Cosmos Token化器。具体来说，采用因果时间卷积层和因果时间注意层来保留视频帧的自然时间顺序，确保使用单一统一的网络架构无缝标记化图像和视频。
直接在高分辨率图像和长视频上训练token化器，而不限制类别或宽高比。与专注于特定数据类别和大小的现有token化器不同，Cosmos Token化器可在各种宽高比下运行 - 包括 1:1，3:4，4:3，9:16 和 16:9。它们在推理过程中与时间长度无关，能够超出训练的时间长度进行token化。
还在标准图像和视频基准数据集上评估 token化器，包括 MS-COCO 2017（Lin，2014）、ImageNet-1K（Deng，2009）和 DAVIS（Perazzi，2016）。为了促进物理 AI 应用的视频token化研究，整理一个视频数据集，该数据集涵盖物理 AI 的许多视频类别，包括鱼眼、机器人、驾驶、人类活动和空间导航。
架构
Cosmos Token化器 被设计为编码器-解码器架构。该架构采用时间因果设计，确保每个阶段仅处理当前和过去的帧，与未来帧无关。与常见方法不同，其在小波空间中运行，其中输入首先由 2 级小波变换处理。具体而言，小波变换以分组方式映射输入视频 𝑥/0:𝑇，以沿 𝑥、𝑦 和 𝑡 将输入下采样四倍。这些组的形成方式为：{𝑥/0、𝑥/1:4、𝑥/5:8、...、𝑥/(𝑇−3):𝑇 } → {𝑔/0、𝑔/1、𝑔/2、...、𝑔/𝑇/4}。后续编码器阶段以时间因果方式处理帧，即 {𝑔/0, 𝑔/0:1, 𝑔/0:2, ...} → {𝜉/0, 𝜉/1, 𝜉/2, ...}。后续编码器阶段遵循类似的方案，最终输出token 𝑧/0:𝑇′。因果设计有助于使在token化器之上构建的模型适应通常在时间因果设置上运行的下游物理 AI 应用程序。小波变换能够在更紧凑的视频表示上进行操作，从而消除像素信息中的冗余，从而使其余层可以专注于更多的语义压缩。
编码器阶段（后小波变换）是使用一系列与下采样块交错的残差块来实现的。在每个块中，采用时空分解的 3D 卷积，其中首先应用核大小为 1 × 𝑘 × 𝑘 的 2D 卷积来捕获空间信息，然后应用核大小为 𝑘 × 1 × 1 的时间卷积来捕获时间动态。用 𝑘 − 1 的左填充来确保因果关系。为了捕获长距离依赖关系，利用具有全局支持区域的时空分解因果自注意 — 例如，最后一个编码器块的 1 + 𝑇 ′。用 Swish 激活函数（Ramachandran，2017）来实现非线性。利用层归一化 LayerNorm (Lei Ba，2016) 代替组归一化 GroupNorm (Wu & He，2018)，这可以防止在潜空间或重建输出的特定区域出现较大的幅度 (Karras，2020；Sadat，2024)。解码器，镜像编码器，用上采样块替换下采样块。如图描绘整个 Cosmos Token化器架构的概览：
采用原始自动编码器 (AE) 公式来建模连续token化器的潜空间。对于离散token化器，采用有限标量量化 (FSQ) (Mentzer，2023) 作为潜空间量化器。连续token化器的潜维度为 16，而离散token化器的潜维度为 6，这表示 FSQ 级别的数量，即 (8, 8, 8, 5, 5, 5)。此配置对应的词汇量为 64,000。
训练策略
采用联合训练策略，以预设频率交替使用小批量图像和视频。仅监督token化器解码器的最终输出。不使用潜空间中的辅助损失，例如承诺或 KL 先验损失。例如，如果将 VAE（Kingma，2013）公式用于连续token化器而不是 原始 AE，则需要有 KL 先验损失。如果将 VQ-VAE（van den Oord，2017）用于离散量化而不是 FSQ，则需要有承诺损失。
采用两阶段训练方案。在第一阶段，用 L1 损失进行优化，以最小化输入和重建视频（𝑥ˆ/0:𝑇）之间的像素级 RGB 差异，以及基于 VGG 特征的感知损失。在第二阶段，使用光流 (OF) 损失 (Teed & Deng, 2020) 来处理重建视频的时间平滑度，并使用 Gram-matrix (GM) 损失 (Gatys et al., 2016) 来增强重建图像的清晰度。此外，在微调阶段使用对抗性损失来进一步增强重建细节，特别是在较大的压缩率下。
以两种压缩率训练图像token化器（表示为 CI 和 DI）：8 × 8 和 16 × 16。同样，以三种压缩率训练视频token化器（表示为 CV 和 DV）：4 × 8 × 8，8 × 8 × 8 和 8 × 16 × 16。这里，图像的压缩率表示为 𝐻 × 𝑊，视频的压缩率表示为 𝑇 × 𝐻 × 𝑊，其中 𝑇 表示时间维度，𝐻 和 𝑊 表示空间维度。
对于视频token化器，创建两个变体：
这种方法确保能够灵活处理图像和视频数据的不同时间和空间分辨率。
预训练 WFM 是通才，可以捕捉现实世界物理和自然行为的一般知识。利用两种不同的可扩展深度学习范式，即扩散模型和自回归模型，来构建两类 WFM。扩散模型和自回归模型，都将困难的生成问题分解为一系列较容易的子问题，并一直在加速生成模型的发展。在扩散模型的情况下，困难的生成问题被分解为一系列去噪问题。在自回归模型的情况下，困难的生成问题被分解为一系列下一个token预测问题。
对于
基于扩散的 WFM
系列，首先分别构建两个 Text2World 模型（7B 和 14B），它们可渲染 Cosmos-1.0-Diffusion-7B-Text2World 和 Cosmos-1.0-Diffusion-14B-Text2World。这些模型可以将文本提示映射到视觉世界的视频。然后，对 Text2World 模型进行微调，以获取额外的视频输入，代表当前观察。结果是一个 Video2World 模型，其中根据当前观察（输入视频）和扰动（文本提示）预测未来视频。这些扩散模型是采用连续token的潜扩散模型。用 Cosmos-1.0-Tokenizer-CV8x8x8 来生成视觉tokens。WFM 的训练文本提示由 VLM 通过视频描述生成生成。这些描述遵循人类对视频描述的不同分布。为了缩小域差距，基于 Mistral-NeMo-12B-Instruct 模型（Mistral & NVIDIA，2024）构建 Cosmos-1.0-PromptUpsampler-12B-Text2World，以帮助将人类文本提示转换为基于扩散 WFM 喜欢的提示。
对于
基于自回归的 WFM
系列，首先构建两个大小分别为 4B 和 12B 的基础模型，以纯粹基于当前视频观察来预测未来视频。分别将它们命名为 Cosmos-1.0-Autoregressive-4B 和 Cosmos-1.0-Autoregressive-12B。这些是从头开始训练的 Llama3 风格 GPT 模型，用于视频预测任务，不具备语言理解能力。为了使基于自回归的 WFM 能够利用文本信息进行下一个token预测，添加到 Transformer 块的交叉注意层，将输入文本提示的 T5 嵌入合并到 WFM 中。这些自回归 WFM 使用 Cosmos-1.0-Tokenizer-DV8x16x16，将输入视频映射到几个整数。Token化器的大量压缩有时会导致不必要的失真。为了解决这个问题，通过微调 Cosmos-1.0-Diffusion-7B-Text2World 模型构建一个扩散解码器（Cosmos-1.0-Diffusion-7B-Decoder-DV8x16x16ToCV8x8x8），将 DV8x16x16 空间中的离散 token 映射到 CV8x8x8 空间中的连续 token。
基于扩散的WFM
基于扩散的 WFM 是潜扩散模型，在token化器学习的潜空间内运行，从而实现视频的紧凑、低维表示。这种设计选择有几个优点：它降低训练和推理过程中的计算成本，同时简化去噪任务（Hoogeboom，2024；Rombach，2022）。为了将视频token化为潜表示，使用 Cosmos-1.0-Tokenizer-CV8x8x8。
为了训练扩散 WFM，采用 EDM（Karras，2022、2024）中概述的方法。与采用高斯流匹配公式（Kong，2024；Polyak，2024）的最新视频生成模型相比，该工作源自扩散分数匹配（Ho，2020；Song，2020）。但是，正如 （Gao 2024）所表明的那样，这些框架在理论上是等效的，在目标和训练程序方面具有根本的相似性。基于 EDM 的公式与这些见解一致，主要区别在于预处理设计和超参数的选择。在实践中，没有遇到 EDM 公式的任何性能限制。
降噪网络 𝐷/𝜃 的设计基于 DiT（Peebles & Xie，2023），最初设计用于标签条件图像生成。本文调整其架构以更好地适应可控视频生成的目标。如图可视化整体网络设计：
3D patch化
。网络的输入是图像和视频数据形状为 𝑇 × 𝐶 × 𝐻 × 𝑊 的潜表示，其中图像通过单帧视频进行区分。为了给降噪网络准备输入，首先使用线性层“patch化”状态，然后将其展平。此过程将形状为 (𝑝/𝑡, 𝑝/h, 𝑝/𝑤) 的非重叠立方体投影到网络的各个token输入中。因此，在patch化之后，图像或视频被重塑为长度为 𝑇𝐻𝑊/(𝑝𝑡𝑝h𝑝𝑤) 的一维时空序列。对降噪网络使用 𝑝𝑡 = 1, 𝑝h = 𝑝𝑤 = 2。
带 FPS-觉察 3D RoPE 的混合位置嵌入和可学习嵌入
。采用一个 3D 分解 RoPE (Su et al., 2024) 来生成任意大小、宽高比和视频长度。具体来说，将特征维度划分为三个大致相等的块，每个块分别沿时间轴、高度轴和宽度轴应用 RoPE 和位置信息。实践中，在各自的轴上连接频率嵌入并重用针对大语言模型 (LLM) 优化的 RoPE 内核，可以有效地实现这一点，无需在每个块中进行拆分和连接。为了进一步支持具有不同帧速率的视频合成，根据训练视频 FPS 重新调整时间频率。由于 RoPE 的相对位置编码属性和 3D 分解设计，FPS-觉察设计与联合图像视频训练兼容。当改变分辨率或视频长度时，RoPE 的另一个好处在渐进训练过程中显而易见。通过利用NTK-RoPE (Peng & Quesnelle，2023)，模型快速收敛，即使在 5,000 个训练步骤内也能实现合理的性能。此外，每个 Transformer 块添加一个额外的可学习绝对位置嵌入可以进一步增强模型，减少训练损失，并减少生成视频中的变形伪影（morphing artifacts）。
用于文本调节的交叉注意
。依靠网络中的交叉注意层来整合语言信息。每个 Transformer 块由连续的自注意、交叉注意和前馈层组成。虽然自注意对时空tokens进行操作，但交叉注意使用 T5-XXL (Raffel，2020) 嵌入作为K-V来集成语义上下文，从而实现有效的文本调节。
Q- K规范化
。在训练的早期阶段，注意 logits 的增长不稳定，导致注意熵崩溃。遵循现有文献 (Dehghani，2023；Esser，2024；Wortsman，2023) 在注意操作之前对 𝑄 和 𝐾 进行归一化。对网络中的所有自注意和交叉注意层使用具有可学习尺度的均方根归一化 (RMSNorm) (Zhang & Sennrich，2019)。
AdaLN-LoRA
。DiT 的 AdaLN 层 (Peebles & Xie，2023；Xu，2019) 占模型参数的很大一部分，而对 FLOPs 方面的计算复杂度的贡献却微不足道。受 W.A.L.T（Gupta，2024）的启发，实施 LoRA（Hu，2022），将这些层中的密集线性投影分解为低秩近似值。对于 Cosmos-1.0-Diffusion-7B，此架构优化实现参数数量减少 36%（从 11B 到 7B 个参数），同时在所有评估指标上保持性能均等，证明参数高效设计的有效性。
联合图像和视频训练
。为了在模型训练中利用大量高质量、多样化的图像数据集，实施一种交替优化策略，将图像和视频数据批次交错。为了促进图像和视频域之间的跨模态知识迁移，采用一种域特定的规范化方案，该方案使用独立估计的图像和视频数据的足够统计数据来对齐潜分布。这种方法的动机是，减少图像和视频潜表示之间的分布漂移可以提高生成质量。此外，在视频潜表示中观察到跨时间和通道维度的非平稳统计数据。为了解决这种异质性，采用一种规范化策略，将逐帧和逐通道标准化应用于视频潜表示，有效地鼓励它们更好地近似各向同性的高斯先验分布。
除了跨模态知识迁移之外，规范化方案还提供重要的理论优势：训练期间信噪比的尺度不变性。考虑两个具有不同尺度的零均值潜表示：一个标准化为单位方差，另一个方差为 4。当添加高斯噪声𝒩(0, 𝜎^2) 以实现标准化表示所需的信噪比时，必须将噪声缩放到𝒩(0, 4𝜎^2) 以使非规范化表示保持相同的比率。通过标准化所有潜表示，确保不同尺度上的信噪比一致，即使在训练期间更新底层token化器，也能促进模型自适应。
为了保持计算效率，平衡图像和视频批次大小，以确保 GPU 之间的内存利用率相当。然而，与图像批次损失相比，视频批次去噪损失表现出较慢的收敛速度。其归因于视频帧中固有的时间冗余，这导致视频批次的梯度幅度较小。从多分辨率图像训练的最新进展（Atzmon，2024；Chen，2023；Hoogeboom，2023）中汲取灵感，将视频批次噪声水平缩放为相对于图像批次噪声水平的帧数平方根来解决这种收敛差异。
渐进式训练
。采用渐进式训练策略，初始阶段涉及以 512 像素的分辨率对视频和图像进行训练，使用由 57 帧组成的视频。随后，过渡到 720 像素的目标分辨率，将视频长度增加到 121 帧。在对海量数据进行预训练后，在高质量子集上对模型进行微调，以线性衰减的学习率进行𝒪（10𝑘）次迭代。与 （Dai 2023) 的研究结果一致，还发现微调可以提高生成视频的质量。
多方面训练
。为了适应具有不同纵横比的内容，将数据组织成五个不同的存储 buckets，分别对应 1:1、3:4、4:3、9:16 和 16:9 的比例，将每个图像或视频分配到纵横比最接近的存储 buckets 中。在训练期间，每个数据并行进程组从一个存储 bucket 中采样，允许不同的并行进程组使用不同的存储 bucket。实施最长边调整大小，以最大限度地保留提示中描述的原始内容信息。对于批处理，对缺失像素应用反射填充（reflection padding），并将填充掩码（padding mask）提供给扩散主干，从而实现推理过程中的精确控制。
混合精度训练
。维护模型权重的两个副本：一个在 BF16 中，另一个在 FP32 中。在前向和后向传递过程中，BF16 权重用于提高训练效率，从而产生同样采用 BF16 格式的梯度和激活。对于参数更新，权重在 FP32 中更新以确保数值稳定性。然后复制更新后的 FP32 参数，并将其转换为 BF16 以进行下一次迭代。为了进一步稳定训练，将去噪分数匹配损失缩小 10 倍。 AdamW 中较低的 betas 和 eps 系数可显著减少损失峰值。对于 14B 扩散模型训练，很少遇到损失峰值，并且没有不可恢复的损失峰值。
文本条件
。对于 Text2World 模型，使用 T5-XXL（Raffel，2020）作为文本编码器。对 T5 嵌入进行零填充以保持固定的 512 序列长度。为了增强文本上下文对齐，采用无分类器指导（Ho & Salimans，2022）。与之前随机将文本嵌入清零的工作（Balaji，2022；Saharia，2022）不同，省略此步骤，因为推理过程中负面提示很有效。值得注意的是，作为文本-到-图像生成器，该模型即使在没有指导的情况下也能生成高保真图像，这种能力归功于高质量的训练数据集。虽然无分类器指导通常会促进对首选视觉内容的模式寻求行为，但仔细的数据选择可以实现类似的效果。但是，对于视频生成，缺乏可比的高质量数据，会导致在低指导设置下结果不理想。因此，在视频生成任务中，需要更高的指导值才能生成令人满意的内容。
图像和视频调节
。扩展 Text2World 模型，构建支持图像和视频调节的 Video2World 模型，方法是将先前的帧合并到生成过程中。具体而言，条件帧沿时间维度与生成的帧连接在一起。为了提高对推理过程中输入帧变化的鲁棒性，在训练期间向条件帧引入增强噪声。此增强噪声的 sigma 值以 𝑃mean = −3.0、𝑃std = 2.0 进行采样。此外，扩散模型的输入，沿通道维度与二进制掩码连接，该掩码将条件帧与生成的帧区分开来。损失函数排除条件帧位置的贡献，只关注生成的输出。为了提高泛化能力，在训练期间随机改变条件帧的数量。在推理过程中，模型可以灵活地使用单个条件帧（图像）或多个先前的帧作为输入。
内存要求
。消耗 GPU 内存的四个主要组件是：
• 模型参数：每个参数 10 字节。混合精度训练将模型参数存储在 FP32 和 BF16 中，同时将指数移动平均 (EMA) 权重存储在 FP32 中。
• 梯度：每个参数 2 字节。将梯度存储在 BF16 中。
• 优化器状态：每个参数 8 字节。使用 AdamW (Loshchilov & Hutter，2019) 作为优化器
并将优化器状态（即一阶和二阶矩）存储在 FP32 中。
• 激活：(2×number_of_layers×15×seq_len×batch_size×d_model) 字节。将激活存储在 BF16 中。为了优化内存使用，实施选择性激活检查点 (Chen，2016；Korthikanti，2023)，重新计算内存受限层（例如规范化函数）的激活。
例如， 14B 模型 (Cosmos-1.0-Diffusion-14B-Text2World) 需要大约 280 GB 用于模型参数、梯度和优化器状态，以及 310 GB 用于高分辨率预训练期间的激活。考虑到 NVIDIA H100 GPU 的 80GB HBM3 限制，采用完全分片数据并行 (FSDP) 和上下文并行 (CP) 在多个 GPU 之间分配内存需求。
完全分片数据并行 (FSDP)
。FSDP 通过在设备之间分片模型参数、梯度和优化器状态来提高内存效率。它仅在计算过程中需要时收集参数，然后释放它们。与跨设备复制参数的标准数据并行不同，FSDP 会分配参数、梯度和优化器状态，每个设备仅管理其分片。此方法将内存使用量最小化到最大的临时未分片参数集及其参数、梯度和优化器状态分片。在实现中，对 7B 模型使用 32 分片因子，对 14B 模型使用 64 分片因子，以平衡内存和通信延迟。
上下文并行 (CP)
。扩展长上下文设置的 transformer 会带来 FLOPs 和激活内存增加的挑战。CP 通过将计算和激活分布在多个 GPU 上来解决这些挑战。它的工作原理是将𝑄 和 (𝐾, 𝑉 ) 沿其序列维度拆分为 CP_SIZE 块，其中 CP_SIZE 是 CP 组内的 GPU 数量。每个 GPU 处理一个 𝑄 块，并使用存储在同一 CP 组中的 (𝐾, 𝑉 ) 块迭代地累积部分注意输出。CP 的不同实现使用不同的通信原语，包括 all-gather (Dubey，2024)、P2P (Liu，2023) 和 all-to-all (Jacobs 等人，2023)。采用 TransformerEngine (NVIDIA, 2024) 的 P2P 变体，它通过在 GPU 之间传输 (𝐾, 𝑉 ) 块同时处理注意来重叠计算和通信。当仔细选择块大小时，这种重叠可以有效地隐藏数据传输延迟。在 NVLink 连接的 GPU 内组织 CP 组，并将 CP 等级与 FSDP 等级重叠以实现最佳利用率。对于具有较短上下文的图像迭代，CP 被禁用以提高吞吐量。由于 (𝐾, 𝑉 ) 的序列长度较短，因此交叉注意层不使用 CP，这导致计算量不足以掩盖通信延迟。
以 Cosmos-1.0-Diffusion-14B 为例，采用分片因子为 64 的 FSDP 可降低参数、梯度和优化器状态的内存要求，使其从 280 GB 降至每 GPU 约 280 / 64 ≈ 4 GB。同样，采用 CP_SIZE = 8 的 CP 可将激活内存从 310 GB 减少到每 GPU 约 310 / 8 ≈ 40 GB。值得注意的是，这些计算是低估的；实际上，token化器和未分片的参数会消耗额外的内存。CP 中的重叠通信和计算还需要每个 GPU 保留多个 (𝐾, 𝑉 ) 块。
在训练过程中，WFM 使用详细的视频描述作为输入文本提示来制作高质量的视频。然而，在推理过程中，用户提示的长度、结构和风格可能会有所不同，通常会短得多。为了弥合训练和推理文本提示之间的差距，开发一个提示上采样器，将原始输入提示转换为更详细和丰富的版本。它可以通过添加更多细节并保持一致的描述结构来改进提示，从而获得更高质量的输出。
提示上采样器的主要要求是：
• 对输入提示的保真度：上采样提示，必须忠实地保留原始用户输入的关键元素，包括主要角色、动作或动作、关键属性和整体意图。
• 与训练分布一致：上采样提示，在长度、语言结构和风格方面，应与 WFM 的训练提示分布非常相似。
• 增强的视觉细节：上采样提示，应设计为提示 WFM 生成更准确的图像。
Text2World 模型的提示上采样器
。对 Mistral-NeMo-12B-Instruct（Mistral 和 NVIDIA，2024）进行微调，以构建提示上采样器。为了获得配对数据，即模拟用户输入的短提示和反映训练提示分布的相应长提示，使用 VLM 根据训练长提示和相应的视频生成短字幕。这种长到短的数据创建策略，可以有效地 (1) 保留来自 WFM 详细训练提示的真实视频内容和分布，以及 (2) 确保短提示和长提示之间的保真度。由此产生的提示上采样器，称为 Cosmos-1.0-PromptUpsampler-12B-Text2World。
Video2World 模型的提示上采样器
。对于 Video2World 模型，输入由视频条件和用户文本提示组成。为了增强用户提示，利用开源 VLM Pixtral-12B（Agrawal，2024）结合零样本提示工程，将提示上采样为详细描述，同时考虑视频条件和用户提示。原始 Pixtral-12B 模型开箱即用，效果很好，因此没有继续执行上述类似的微调。
基于自回归的WFM
在自回归 WFM 中，将世界模拟生成公式化为类似于语言建模的下一个token预测任务。首先使用 Cosmos 离散token化器将视频转换为离散视频token序列 𝒱 = {𝑣/1, 𝑣/2, . . . , 𝑣/𝑛}。然后训练 Transformer 解码器 (Vaswani et al., 2017)，使用过去的视频tokens作为上下文来预测下一个视频token，类似于大语言模型 (LLM) (Brown et al., 2020; Dubey et al., 2024; Jiang et al., 2023)。具体来说，训练目标是最小化一个根据下一个视频 token 预测的负对数似然 (NLL) 损失。
基于自回归的 WFM 架构如图所示。对针对视频生成任务量身定制的标准 Transformer 模型架构进行几项修改，包括添加 1) 3D-觉察位置嵌入、2) 交叉注意以启用文本输入以实现更好的控制，以及 3) QK- 规范化 (Wortsman，2023)。
3D 位置嵌入
。与基于扩散的 WFM类似，结合两种互补的位置嵌入机制：用于相对位置的 3D 分解 RoPE 和用于绝对坐标的 3D 分解绝对位置嵌入 (APE)。这些机制协同工作，为整个网络提供全面的空间和时间信息。
•
3D RoPE
。将 3D RoPE 应用于模型，以在时间、高度和宽度维度上编码相对位置信息。在训练期间，采用多阶段训练策略，其中视频的序列长度随训练的进行而增加。为了使 3D RoPE 适应不断变化的时间持续时间，使用 YaRN（Peng，2023），这是一种旨在扩展 RoPE 上下文窗口计算效率高的技术。仅沿时间轴应用 YaRN 扩展，因为视频序列长度仅沿时间维度增加。通过利用 YaRN，模型可以推断出比训练初始阶段遇到的上下文长度更长的上下文长度。
•
3D APE
。除了 3D RoPE，还在每个 Transformer 块中加入 3D APE 来补充相对位置编码。此 APE 使用跨时间、高度和宽度维度分解的正弦嵌入对位置信息进行编码，确保模型知道绝对位置。嵌入在每个阶段直接添加到输入张量中，丰富 Transformer 的位置上下文。结合绝对和相对位置编码可以提高模型性能、减少训练损失并最大限度地减少生成视频中的变形伪影（morphing artifacts）。值得注意的是，虽然基于扩散的 WFM 采用可学习的嵌入，但在基于自回归的 WFM 中为 APE 采用基于正弦的嵌入。
词汇表
。Token化是将输入文本转换为大语言模型 (LLM) 中的离散token序列的关键步骤。在 LLM 中，可能的token 词汇表由 LLM 的token化器（例如，OpenAI (2022) 引入的 tiktoken）确定，该token化器使用字节对编码 (BPE)（Gage，1994）等算法在大量文本上进行训练。
对于自回归模型，使用 Cosmos-1.0-Tokenizer-DV8x16x16 作为token化器。利用有限标量量化 (FSQ) (Mentzer，2023) 将 6 维潜空间量化为 (8, 8, 8, 5, 5, 5) 级别。这种量化导致词汇量大小为 8×8×8×5×5×5 = 64,000。
