Cosmos World Foundation Model Platform for Physical AI

Introduction

研究贡献

本文提出了一个新的平台——Cosmos世界基础模型平台（Cosmos World Foundation Model Platform），旨在为物理AI开发者创建定制化的世界模型。其核心思想是利用视觉世界基础模型（以视频形式呈现观察数据）通过预训练和后训练范式解决数据扩展难题：

1. 预训练世界模型（Pre-trained WFM）： 利用大规模视频数据集构建通用模型，涵盖丰富的视觉体验。
2. 后训练世界模型（Post-trained WFM）： 使用目标物理AI环境的数据集对预训练模型进行微调，从而适配具体任务。

平台特点

Cosmos平台提供了多种关键组件：

1. 视频数据筛选管道： 从20M小时的视频中提取1亿高质量动态片段。
2. 视频标记与编码器： 开发了一种因果设计的视频标记方法，压缩视频信息以提高训练效率。
3. 预训练世界模型：two scalable approaches for building pre-trained world foundation models—the diffusion model and the autoregressive model
4. 后训练的示例： 包括相机控制、机器人操控和自动驾驶等应用场景。
5. 开放性与共享性： 提供开放源码和模型权重，降低开发门槛。

Data Curation

数据处理目标

为了构建高质量的训练数据集，文章设计了一条完整的视频数据筛选与处理管道。这条管道旨在从庞大的原始视频集合中提取有助于世界基础模型（WFM）训练的高质量数据，同时保证数据多样性和高效性。

数据来源

1. 数据总量：使用了20M小时的视频（分辨率从720p到4k不等），包括专有视频数据集和公开的网络视频。
2. 目标覆盖范围：数据覆盖多个物理AI应用场景和类别，例如驾驶（11%）、手部操作与物体操控（16%）、人类运动与活动（10%）、空间导航（16%）、自然动态（20%）等。
3. 挑战：视频数据通常存在以下问题：
   • 编码格式和设置的多样性；
   • 长度、分辨率和场景过渡的变化；
   • 视频可能含有不必要的后期处理或编辑效果，影响模型训练质量。

数据处理流程

如图所示，数据处理流程包含以下五个步骤：

1. 分割（Splitting）

• 目标：将长视频分割为无场景切换的短片段，并保证每段的视觉一致性。
• 操作步骤：
  1. 检测视频中的镜头切换点，分割视频为不超过60秒的片段；
  2. 丢弃短于2秒的片段（可能是切换或视觉特效）。
• 技术方法：评估了多种镜头分割算法（如PySceneDetect、Panda70M、TransNetV2、AutoShot），最终选择TransNetV2因其对复杂场景切换的处理效果较好。
• 重新编码：将分割的视频重新编码为高质量的mp4格式，以提升后续处理的效率和一致性。

2. 过滤（Filtering）

对分割后的片段进行过滤，移除低质量或无意义的视频内容。包括以下四个子步骤：

1. 运动过滤：
   • 去除静态视频或有随机剧烈镜头运动的片段；
   • 使用NVIDIA加速的光流估计网络，结合轻量级分类器进行动作类型标记（如平移、缩放、倾斜等）。
2. 视觉质量过滤：
   • 检测并移除含有失真（如模糊、噪点、过度曝光、欠曝光等）的片段；
   • 使用视频质量评估模型（如DOVER）和图像美学模型（Schuhmann, 2022）评分，去除低于质量阈值的视频。
3. 文本覆盖过滤：
   • 检测后期处理中添加了大量文字或特效的视频；
   • 训练了基于MLP的分类器，利用视频嵌入（如InternVideo2）进行文字检测。
4. 视频类型过滤：
   • 针对物理AI任务设计视频内容分类（如人类动作、物体交互等）；
   • 通过分类器标注片段，并上采样关键类别、下采样不相关类别。

3. 标注（Annotation）

• 使用视觉语言模型（VLM）为每段视频生成高质量、细粒度的描述文本。
• 测试了多种VLM（如VILA、VFC、Qwen2-VL），最终选择性能最优的VILA模型进行标注。
• 提高效率：利用TensorRT-LLM引擎对VILA模型进行量化，实现了10倍的推理速度提升。

4. 去重（Deduplication）

• 目标：减少语义冗余，提高数据多样性并降低训练复杂度。
• 技术方法：基于SemDeDup和DataComp的语义去重方法，使用InternVideo2嵌入和加速k-means聚类算法检测近似重复视频。
• 结果：去除了约30%的重复数据，同时保留了分辨率最高的视频。

5. 分片（Sharding）

• 将处理后的数据片段打包为模型训练可直接消费的WebDataset格式；
• 根据视频的分辨率、长宽比和时长进行分片，以适应训练的课程设计。

基础设施支持

• 使用AnyScale Ray框架实现流式数据处理，适应分布式集群；
• 提升效率：采用了分段资源优化（如解码、计算和传输并行操作），结合分布式调度算法（Fragmentation Gradient Descent），显著提高了多资源利用率。

结果

最终生成了约1亿个高质量的视频片段（用于预训练），以及约1千万个高精度视频片段（用于微调），为后续模型训练奠定了坚实的数据基础。

Tokenizer

Tokenizer 是现代大规模模型的重要组成部分，其作用是将原始数据（例如图像和视频）转化为紧凑的语义表示，以便提高训练和推理的效率。本文提出了 Cosmos Tokenizer，支持图像和视频的离散与连续标记，具有优异的压缩质量与高效推理能力。

Cosmos Tokenizer 的设计特点

多功能性

• 支持因果性（Causal Design）：
  视频标记是因果的，即当前帧的标记仅依赖于过去帧的观察，这与物理AI系统的因果性质对齐。
• 联合标记（Joint Tokenization）：
  Cosmos Tokenizer 能够同时处理图像和视频数据，允许将图像数据用于视频模型的训练，提高模型的视觉表现力。
• 支持图像与视频的统一处理
  视频是由时间维度上的多个帧组成，而图像可以被看作只有单帧的“视频”。标记过程只依赖于当前帧及之前的帧，保证处理的时间因果性。因此，对于图像（单帧），标记过程与视频一致，只处理当前帧即可。
• 训练策略支持图像与视频的联合处理
  在训练阶段，Cosmos Tokenizer 会交替输入图像和视频

压缩率与多样性

• 支持不同的分辨率（如1:1、16:9）和视频时长的标记。
• 在标记时对输入进行小波变换（Wavelet Transform），减少像素冗余信息，保留高语义信息。

架构设计

Cosmos Tokenizer 采用了轻量级的编码器-解码器（Encoder-Decoder）架构，其特点如下：

1. 因果时间卷积与注意力机制：
   在时间维度上，采用因果卷积与因果注意力层，以保持标记的时间因果性。
2. 小波变换预处理：
   输入视频先经过小波变换，减少输入的维度，压缩3D视频的空间和时间冗余信息。
3. 残差与下采样模块：
   编码器使用残差块和降采样模块捕获时空特征。解码器则通过上采样块还原输入。

训练目标

1. 连续标记：
   使用自编码器（Autoencoder，AE）学习连续的潜在空间。
2. 离散标记：
   使用有限标量量化（Finite-Scalar-Quantization, FSQ）技术学习离散的潜在空间。
3. 损失函数：
   • 像素损失（L1）： 最小化输入和重建视频之间的像素差。
   • 感知损失（Perceptual Loss）： 基于VGG-19特征的感知差异。
   • Flow损失（Flow Loss）： 提高视频重建的时间平滑性。
   • Gram矩阵损失（Gram Loss）： 提升重建图像的清晰度。

性能评估

压缩率 vs. 质量

• Cosmos Tokenizer 在连续与离散标记上都实现了较高的重建质量和较好的压缩率。
• 相比现有方法，Cosmos Tokenizer 在DAVIS视频和标准图像数据集（如ImageNet-1K）上显著提升了峰值信噪比（PSNR）和结构相似性指数（SSIM）

推理效率

• 在单张 NVIDIA A100 GPU 上，Cosmos Tokenizer 比现有方法快 2倍到12倍，且模型参数更少，计算效率更高。