AI Research 技术调研

WorldDreamer: Towards General World Models for Video Generation via Predicting Masked Tokens

13分钟阅读

一句话定位

WorldDreamer 是 GigaAI 与清华大学团队 2024 年 1 月提出的通用世界模型：借鉴 LLM 的 masked-language-modeling 思路，把图像/视频编码成 VQGAN 离散 token 后随机 mask 一部分，用一个时空局部注意力的 Transformer（STPT）预测被 mask 的 token，并通过 cross-attention 注入文本（T5）与动作（MLP）双模态条件；单一模型同时支持图生视频、文生视频、视频修补、视频风格化、动作到视频（驾驶场景）五类任务，并靠 parallel mask-token 解码（约 10 步）声称比扩散/自回归方法快 3×~20×。论文只有定性可视化结果，无任何量化 benchmark 表格，GitHub 仓库自建立以来也从未发布代码或权重。

背景与定位

在此之前，世界模型的研究基本被限定在单一场景内：游戏场景的 World Models（Ha & Schmidhuber）与 Dreamer 系列、机器人场景的 daydreamer，以及自动驾驶场景的 gaia-1-wayve、drivedreamer、ADriver-I 等。这些方法要么依赖 RNN 隐状态（Dreamer 系列），要么依赖扩散主干（GAIA-1 的解码器、DriveDreamer），普遍缺乏跨场景的通用性。

WorldDreamer 的定位是做第一个跨自然场景与驾驶场景的通用世界模型，并在建模范式上做了明确选择：不走 GAIA-1/VideoPoet（videopoet）式的自回归 next-token 预测，也不走扩散去噪，而是把 maskgit / microsoft-wham-muse 的 mask-and-predict 并行解码范式首次系统性地搬到世界模型构建上——即把视觉动态学习 cast 成”无监督视觉序列建模”问题，用未被 mask 的 token 预测被 mask 的 token（类比 BERT 之于 NLP）。同时提出 Spatial Temporal Patchwise Transformer（STPT），在 U-ViT 基础上把 self-attention 限制在时空局部 patch 窗口内，以加速收敛、降低计算量，再用 cross-attention 把文本/动作多模态 prompt 接入。范式命名：masked visual token prediction world model（LLM 式非自回归并行解码世界模型）。

模型架构

  • 视觉分词器：VQGAN，将 $N$ 帧视觉输入 $I\in\mathcal{R}^{N\times H\times W\times 3}$ 编码为离散 token $T_V\in\mathcal{R}^{N\times h\times w}$，vocabulary size 为 8192，用数十亿张图像（LAION-5B）预训练。论文正文对下采样倍率的两处表述不一致（原文如此）：叙述称”downsamples the resolution by 16ד，但紧接给出的公式为 $h=H/4, w=W/4$（即 4× 下采样），未见勘误说明，此处照录两处原文数字。
  • STPT 主干：基于 U-ViT 改造的 Spatial-Temporal Patchwise Transformer。输入 token 先通过一个可学习 codebook（size 8193 = VQGAN 的 8192 + 1 个 mask token）映射为视觉 embedding $E_V\in\mathcal{R}^{N\times h\times w\times C_V}$；每层内先过 3D 卷积，再沿空间维切成 patch（patch stride $s=2$，经验设定），patch 内做展平后的自注意力（spatial-temporal patchwise self-attention），随后 reshape 回原维度；接着做空间维 cross-attention，把帧数当作 batch 维度，将多模态 embedding $E_M$ 融入视觉特征。$L$ 层 STPT 后，特征维度映射到 VQGAN 的 codebook size，用 softmax 计算每个被 mask token 的概率分布。
  • 文本条件：预训练 T5 把文本编码为 $E_T\in\mathcal{R}^{K\times C_T}$（$K$ 为序列长度），repeat 到 $N$ 帧并映射到通道数 $C_V$。
  • 动作条件：一个 MLP 把动作输入（NuScenes 场景下为 ego 车的 yaw 角与速度）编码为 $E_A\in\mathcal{R}^{N\times C_V}$。
  • 文本、动作 embedding 拼接为多模态 prompt $E_M\in\mathcal{R}^{N\times(K+1)\times C_V}$；两者均可为空以支持无条件/单条件生成。
  • 图像-视频联合训练：图像输入时把自注意力 $\mathcal{F}_s$ 的 attention weight 替换为对角矩阵，使 STPT 可与视频数据混合训练，扩充训练样本规模。
  • 模型规模：finetune 阶段 STPT 全部 1B 参数均可训练（论文未进一步拆分层数/隐藏维等具体配置）。
  • Mask 策略：训练时按 cosine schedule 采样随机 mask 率 $r\in[0,1]$，每帧共 mask $\frac{2hw}{\pi}(1-r^2)^{-1/2}$ 个 token；同一时空窗口内不同帧使用相同的token mask（避免相邻帧相似内容导致的信息泄漏）。推理时同样遵循 cosine mask schedule，每步选置信度最高的一部分 token 解 mask 并固定，其余继续预测，约 10 步完成全部生成。

数据

  • Deduplicated LAION-2B：对 LAION-5B 做去重、过滤水印概率 >50% 或 NSFW 概率 >45% 的图像后得到的精简子集，用于图像端预训练。
  • WebVid-10M：约 1000 万条短视频，平均时长 18 秒，主要分辨率 336×596，每条配文本描述；因视频带水印会污染生成结果，仅用于早期预训练阶段。
  • 自采集视频-文本对：从公开网络获取视频，用 PySceneDetect 做转场检测切出单一连续场景片段，再用光流幅度过滤慢动作片段，最终得到 50 万条高质量视频片段；取每条视频第 10/50/90 百分位帧作为关键帧，用 Gemini 分别生成关键帧 caption 并聚合为整段视频的详细 caption。
  • NuScenes：700 条训练视频 + 150 条验证视频，每段约 20 秒 @ 12Hz；仅用前视相机视频，采样间隔为 6 帧，最终得到约 2.8 万条驾驶场景训练视频；用 Gemini 生成含天气、时段、道路结构、关键交通元素的逐帧描述并聚合为整段 caption；同时提取 ego 车的 yaw 角与速度作为动作元数据。
  • 训练阶段划分：先在 WebVid + LAION 混合数据上预训练（WebVid 每条取 16 帧，LAION 每 16 张独立图像组成一个训练样本，输入分辨率 256×256），再在自采集数据 + NuScenes 上 finetune（输入分辨率 192×320，每样本 24 帧）。

训练方法

  • 目标函数：交叉熵损失（式 5），最大化在给定未被 mask token 与多模态 prompt 条件下、被 mask token 的对数似然（式 2）——本质是非自回归的并行 mask-token 预测（BERT/MaskGIT/Muse 同族思路）。
  • 预训练：2M iterations，batch size 64，AdamW 优化器，学习率 5×10⁻⁵，weight decay 0.01。
  • Finetune：在自采集数据 + NuScenes 上继续训练 20K iterations，batch size 32，学习率 1×10⁻⁵，STPT 全部 1B 参数可训练。
  • 推理时的 Classifier-Free Guidance：训练阶段随机丢弃 10% 样本的多模态 embedding；推理时对每个被 mask token 分别算条件 logit $c$ 与无条件 logit $u$，按 $g=(1+\beta)c-\beta u$ 合成最终 logit（$\beta$ 为 guidance scale）。
  • 生成流程：cosine mask schedule 迭代式并行解码，约 10 步完成，相比扩散方法通常需要的约 30 步降噪、自回归方法通常需要的约 200 步逐 token 预测，声称有 3×~20× 的步数级加速（引言中另有”较扩散方法快约 3ד的单独表述）。
  • 无强化学习环节，纯自监督/无监督的视觉 token 预测训练；论文未披露具体 loss 曲线或消融实验。

Infra（训练 / 推理工程）

  • 训练所用 GPU 型号、数量、GPU-hours、训练精度（FP16/BF16）、并行策略：论文全文未披露。
  • 推理速度：在单张 A800 上生成一段 24 帧、分辨率 192×320 的视频耗时 3 秒（对应约 10 步 parallel 解码）。
  • 速度对比均为步数级的概念性比较（约 10 步 vs 扩散约 30 步 / 自回归约 200 步），论文未给出具体 wall-clock 基准对比表。

评测 benchmark

论文没有任何量化 benchmark——无 FVD/FID/CLIP score 等指标，也没有与任何 baseline 的定量对比表格。全部为定性可视化实验（§4.3）：

  • 图生视频（Fig.5）：单张图像输入生成后续帧，展示类似真实电影运镜的连贯运动，且严格遵循初始图像的构图约束。
  • 文生视频（Fig.6）：不同风格文本 prompt 驱动生成对应风格/内容/运镜的视频。
  • 视频修补（Fig.7）：给定感兴趣区域 mask + 文本 prompt，对视频局部内容做替换。
  • 视频风格化（Fig.8）：给定随机 token mask + 风格 prompt（如”秋天主题”），对整段视频做风格迁移。
  • 动作到视频（Fig.9，NuScenes 驾驶场景）：同一起始帧配合不同驾驶动作（左转/右转）生成不同的未来帧序列。

论文明确未提供任何数值化的生成质量评测或消融实验表格。

创新点与影响

  • 贡献：(1) 提出首个跨自然场景 + 驾驶场景的通用世界模型，突破此前世界模型局限于游戏/机器人/驾驶单一领域的边界；(2) 提出 STPT，把时空局部 patch self-attention 与空间维 cross-attention 结合用于视觉 token 预测，兼顾效率与多模态条件注入；(3) 把 mask-token parallel 解码（而非扩散去噪或自回归 next-token）引入世界模型构建，声称继承 LLM 训练配方的可扩展性优势，同时相比扩散/自回归方法有显著推理加速。
  • 改变了什么：为”用 MaskGIT/Muse 式并行 mask-token 预测构建世界模型”提供了早期范例，与同期 GAIA-1（自回归）、DriveDreamer（扩散）的路线形成对照，是该并行解码范式向”通用世界模型”方向的一次系统性尝试。
  • 局限（论文正文未设专门局限章节，以下为可核实的事实性局限）：全文不含任何量化评测或消融实验，无法验证生成质量、时序一致性、物理合理性相对 baseline 的优劣；VQGAN 下采样倍率的两处表述（16× vs $h=H/4$）在原文中互相矛盾；GitHub 仓库自 2024-01-18/19 建立后只完成”Repository Initialization”提交（GitHub API 显示 pushed_at 停留在 2024-01-19），从未发布代码、训练脚本或模型权重，无法复现。

原始链接

一手源存档（sources/）

关系图谱

反向链接