AI Research 技术调研

Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

18分钟阅读

一句话定位

Diffusion Forcing（DF）是 MIT CSAIL 2024 年 7 月提出的一种训练与采样范式：让扩散模型对一组 token 各自独立地施加不同噪声级别去噪，从而把”下一 token 预测（teacher forcing）“和”整序列扩散（full-sequence diffusion）“统一在同一个模型里——因果版本 Causal Diffusion Forcing（CDF）既能像自回归模型那样变长、逐 token 生成，又能像整序列扩散那样对整条轨迹做引导（guidance）。它解决了自回归视频生成滚动超出训练长度就发散的顽疾（可稳定 rollout 到 1000+ 甚至 2000 帧），并在 D4RL 规划、真实机器人模仿学习上带来显著增益。这是世界模型”可控性/长时程稳定”方向的一块方法论基石。

背景与定位

DF 站在两类序列生成范式的交叉点，论文用一张能力表（图 1）把三者对齐：

  • 下一 token 预测（teacher forcing [Williams&Zipser 1989]）：LLM 主流范式，优点是变长生成、可条件于任意长历史、支持树搜索与在线反馈控制；缺点是没有引导机制、且在连续数据（视频）上自回归 rollout 极易发散——前一帧的微小误差逐帧累积，模型崩溃。
  • 整序列扩散（full-sequence diffusion）：视频生成（cosmos-predict 一路的 [Video Diffusion Models]）与长时程规划（Diffuser [Janner 2022]）主流，优点是可用 classifier/classifier-free guidance 把采样引向高奖励轨迹、擅长连续信号；但它普遍用非因果、无掩码架构，只能整条固定长度生成，限制了变长与子序列生成，也限制了引导方式。
  • 论文指出：天真地把”下一 token 预测”直接拿去做整序列扩散会生成很差——因为它没建模”早期 token 的小不确定性必然放大成后期 token 的大不确定性”这一因果结构。

DF 的核心 insight 是把”加噪”重新解释为”部分掩码”：噪声为 0 = token 完全可见（unmasked），噪声为 K = token 完全被掩码（纯白噪声）。于是训练模型去”unmask 任意一组被不同程度加噪的 token”。把时间轴（Bayes 滤波，下标 t）与噪声轴（扩散的”不确定性”轴，上标 k）交织起来，就得到一个统一视角。最接近的先验工作是 AR-Diffusion（因果文本扩散，噪声级别沿时间轴线性依赖）与 Rolling Diffusion；DF 的区别是每个 token 噪声级别完全独立，可在采样时用一个 2D 调度矩阵任意编排。范式命名：per-token independent noise level diffusion / Causal Diffusion Forcing。

与同期”神经游戏引擎 / 可玩世界模型”（deepmind-genie2、decart-oasis、diamond）相比，DF 不是一个具体的世界模型产品，而是提供了它们都需要的底层能力——因果、变长、稳定长滚动 + 序列级引导；其后续 history-guidance（Diffusion Forcing Transformer / History-Guided Video Diffusion）把它搬到 DiT + 潜空间扩散上做长视频生成。NeurIPS 2024 收录。

模型架构

总范式：训练一个共享的”下一/下几个 token”预测模型，对序列 (x_t^{k_t})_{1≤t≤T} 去噪，其中每个 token t 的噪声级别 k_t 可独立取值。单个 RNN unit 的输入输出行为等价于一个标准条件扩散模型：以 z_{t-1}（历史 latent）和噪声 token x_t^{k_t} 为条件，预测无噪 x_t=x_t^0（进而通过仿射重参数化间接预测噪声 ε）。

  • 因果骨干（本文极简实现）：vanilla RNN。RNN 维护 latent z_t，按 z_t ~ p_θ(z_t | z_{t-1}, x_t^{k_t}, k_t) 递推。当 k_t=0 这是 Bayes 滤波的后验更新；k_t=K（纯噪声、无信息）时退化为建模先验分布 p_θ(z_t|z_{t-1})。z_t 实现为确定性（代表”信念的分布”而非采样），以便对 latent dynamics 反向传播。作者说明 transformer 实现也可行（附录 B.1），但 RNN 对在线决策更高效，故视频/规划都用 RNN。
  • 视频扩散架构：把原图 x 与 latent z 都取为 2D 张量（channel × width × height，且 x、z 同宽高）。转移模型 p(x_t^{k_t}|z_{t-1}) 用标准扩散 U-Net；U-Net 输出送入 GRU（以 z_{t-1} 作隐状态）得到 z_t；观测模型 p(x_t|z_t) 用 1 层 ResNet + 一个 conv 层。三者组合成一个 RNN 层。
  • 非视频架构：用 残差 MLP（resMLP） 替代 U-Net 作为动力学骨干，另一个 resMLP 作观测模型（同样接 GRU）。
  • 配置数字：latent z 通道数 = DMLab 16 / Minecraft 32；Minecraft 模型 36M 参数、DMLab 模型 24M 参数、maze 规划模型 4.33M 参数；时序模型 = 1 个 MLP + 4 个 GRU。
  • 扩散参数化：视频用 v-parameterization（对收敛速度与画质都关键）；规划/模仿学习用 x0-parameterization（不人为强调高频细节）；时序用 v。
  • 噪声调度：视频 sigmoid、maze 规划 linear、其余 cosine。
  • 采样（Algorithm 2）：定义一个 2D 调度矩阵 K ∈ [K]^{M×T}——列对应时间步 t、行 m 对应噪声级别。全序列初始化为白噪声（k=K），逐行向下、每行内从左到右按 K 规定的噪声级别去噪，最后一行 K_{0,t}≡0 得到干净序列。由于 DF 被训练成能对”任意噪声级别序列”去噪，同一个模型无需重训即可用不同的 K 实现不同行为。

采样带来的新能力（本文核心卖点）：

  • 稳定自回归 rollout：用”上一时刻略带小噪声（0<k≪K）的 latent”来更新，从而在 rollout 超出训练长度时不发散。
  • 保持未来不确定（zig-zag / pyramid 调度）：近未来先充分去噪、远未来保持高噪声，把噪声级别当作”不确定性”来编码”越远越不确定”。
  • 长时程引导：因未来 token 依赖过去，未来 token 的引导梯度可沿时间反向传播去引导过去 token 的采样，同时尊重因果性。
  • Monte Carlo Guidance（MCG）：对 x_t^k 的引导用”整个未来 x_{t+1:T} 分布上的期望奖励”——采样多条未来、平均其引导梯度（类似 MPPI 打靶法），配合 zig-zag 调度效果更强。这是非因果整序列扩散做不到的。

数据

本文是方法论文，横跨 5 个域，各自数据规模具体如下：

  • 视频（Minecraft）：取自 TECO 数据集（github.com/wilson1yan/teco）的 Minecraft gameplay，128×128 像素，裁成 72 帧 子序列；数据自带配对动作但故意丢弃以增加预测随机性；因算力仅用约 10% 子序列训练。
  • 视频（DMLab）：TECO 的 DeepMind Lab 3D 迷宫导航随机游走，64 像素分辨率；训练帧长论文正文/图注写 36 帧、附录 F.1 另处写 48 帧（原文内部略有出入，此处如实标注），同样丢弃动作、仅用约 10% 数据。
  • 规划（D4RL）：maze2d-{umaze,medium,large}-v1 三个环境（含 multi-goal 变体共 6 个评测设置）。观测 4 维（2D 位置 + 速度），动作 2 维（加速度），稀疏奖励（离目标 0.5 半径内得 1）。数据集 = 迷宫内随机游走（在路点间随机导航、避墙），从不指向特定目标（引入随机性）。每个迷宫训一个模型；沿用 Diffuser 约定丢弃数据集奖励、只用目标位置规划。
  • 机器人（真实 Franka）：一个需要记忆的长时程任务——把苹果/橙子用第三个空槽互换位置（当一个水果在中间槽时，单帧观测无法判断下一步，必须记住初始配置，故非马尔可夫）。150 条专家演示，VR 遥操 + 阻抗控制采集；双相机（手腕相机 + 前方全景）；每条 500–600 帧、含 6 自由度 动作；padding + 下采样到 40 帧，每帧捆绑 15 个动作。
  • 时序：GluonTS 的 6 个真实数据集 Exchange / Solar / Electricity / Traffic / Taxi / Wikipedia（多变量高维），context 窗口与 prediction 窗口同长。

工程技巧 · frame stacking：把多个相邻观测堆成一个 x 以省算力/显存并复用扩散模型的标准超参——DMLab 堆 4 帧、Minecraft 堆 8 帧、maze 规划堆 10 帧。

训练方法

  • 目标函数（式 3.1）：常规扩散噪声预测目标。训练时对每条数据 x_{1:T}，从 [K]^T 均匀采样噪声级别序列 k_{1:T}，对每个 token 前向加噪 ε_t ~ N(0, σ²_{k_t} I)，最小化 Σ_t ||ε_t − ε_θ(z_{t-1}, x_t^{k_t}, k_t)||²（Algorithm 1）。
  • 理论保证（Theorem 3.1）：该训练过程优化的是 ln p_θ((x_t^{k_t})) 期望对数似然的 ELBO 的一个 reweighting；在适当条件下同时最大化所有噪声级别序列的似然下界。一个特例是 k_t∈{0,K}——即可掩掉任意历史 token，DF 便能从正确的条件分布采样，等价于建模训练集所有可能子序列的分布。
  • Fused SNR reweighting（附录 D.1）：在 min-SNR 策略 [Hang 2024] 基础上，把”噪声观测自身的 SNR”与”历史累积 SNR（指数衰减 running mean）“按独立事件概率融合成一个 reweighting 因子，显著加速视频预测收敛（非图像域没观察到增益故未用）。
  • 决策框架：token 设为 x_t=[a_t, r_t, o_{t+1}]（动作/奖励/下一观测），一条轨迹即 x_{1:T}，训练同 Algorithm 1。执行时用 latent z_{t-1} 总结历史，采样出计划 x̂_{t:t+H}（H = 前瞻窗口，类 MPC），执行 â_t 后按后验 p_θ(z_t|z_{t-1}, x_t, 0) 更新 latent。同一模型既是 policy（缩短 H 降延迟）又是 planner（拉长 H + 引导做长时程规划），无需改架构或重训。引导可用任意奖励：整轨迹稠密奖励、前瞻窗口稠密奖励、或稀疏 goal 完成 −||o_T−g||²。
  • 推理加速：每个 token 的扩散用 DDIM 采样；训练 K=1000 步扩散，采样时视频只用 100 DDIM 步、非视频 50 步。采样时（Algorithm 2 第 8 行）直接令 z_t ← z_t^new 而非重算后验，算力砍半，且恰好实现稳定化所需行为。

Infra（训练 / 推理工程）

  • 精度：全部实验用 fp16 混合精度训练。
  • 视频预测（GPU 密集）：8× A100，训练 50K 步，batch size 8×16，通常 12 小时 在约 40K 步收敛（含偶尔验证时间），单次训练 < 1 天。
  • 其余域（时序 / maze 规划 / compositional / 视觉模仿）：单张 2080Ti（11GB） 即可。batch size：maze 与 compositional 2048、视觉模仿 32、时序 32；通常 50k–100k 步收敛，4–8 小时。
  • 推理侧：作者明确选 RNN 而非 transformer 作骨干，因其”在线决策的灵活与高效”；论文未给出具体的控制频率 / FPS / 边缘硬件时延数字（未披露）。frame stacking 亦同时降训练时间与显存。

评测 benchmark

均来自论文一手表格/图：

1) 视频预测（定性为主，图 3 / 图 8–13）：DF 在 Minecraft（训练 72 帧）与 DMLab（训练 36 帧）上都能稳定 rollout 到 180 帧乃至 1000、2000 帧不发散（2–5× 甚至”近乎无限”超训练长度）；而两个共用同一 RNN 架构的基线——teacher-forcing 的下一帧扩散、以及因果整序列扩散——很快发散。此外训练长度内，整序列扩散会出现帧间跳变不连续，DF 则表现出穿越一致 3D 环境的自洽 ego-motion。

2) D4RL 规划（Table 1，episode reward，越高越好）：DF 在全部 6 个环境上超过 Diffuser 及所有基线。

设置MPPICQLIQLDiffuser*Diffuser（执行生成动作）Ours w/o MCGOurs
Single-task 平均16.27.747.0119.58.7129.67141.7
Multi-task 平均21.5–16.9129.420.6127.7146.2

关键结论：(a) MCG 有效——去掉 MCG 性能下降（141.7→129.67 等），但 DF 仍具竞争力；(b) 建模因果的价值——Diffuser 生成的 state 与 action 不因果自洽，其实现只能丢弃生成动作、改用手工 PD 控制器（表中带 *），一旦直接执行生成动作性能崩到 8.7 / 20.6，而 DF 的原始动作生成自洽、甚至优于”Diffuser 状态 + 手工 PD”；(c) 灵活 horizon——DF 天然支持随任务推进收缩规划 horizon。

3) 真实机器人模仿（4.4 节）：需要记忆的水果互换任务，DF 成功率 80%，无记忆的 diffusion policy（SOTA 模仿学习）失败。鲁棒性：加视觉干扰甚至完全遮挡相机时，用 k>0 把观测标为”噪声”让模型更依赖先验模型，成功率仅降 4% 到 76%；而下一帧扩散基线必须把扰动当真值、OOD 误差大，只有 48%。同一模型还能仅给单帧就生成机器人执行任务的视频（打通模仿学习与视频生成，指向”用无标注视频预训练”）。

4) 多变量时序预测（Table 2，CRPS_sum，越低越好）：DF 与 TimeGrad / ScoreGrad / Transformer-MAF 等整体持平（Exchange 上 0.003±0.001 最优），仅 Wikipedia 排第四。作者强调时序非核心应用，只为证明新训练目标不会退化其作为通用序列模型的能力。

创新点与影响

  • 核心贡献：提出以”每 token 独立噪声级别 = 部分掩码”统一 next-token prediction 与 full-sequence diffusion，用一个 2D 噪声调度矩阵在采样时自由切换/混合两种行为，无需重训。
  • 它改变了什么：给了世界模型/序列生成三样此前无法兼得的能力——因果 + 变长 + 稳定超长滚动，外加尊重因果的长时程引导与 Monte Carlo Guidance；直接治好了连续域自回归 rollout 的发散病，并让”同一模型既当 policy 又当 planner”。
  • 理论：证明训练目标是”所有子序列联合似然”ELBO 的 reweighting（Theorem 3.1）。
  • 作者自陈的局限：当前因果实现基于 RNN，更高分辨率视频/更复杂分布可能需要大 transformer（附录 B.1）；未研究扩到互联网规模数据/任务的 scaling 行为。
  • 后续影响：官方随即发布 Diffusion Forcing v2 / History-Guided Video Diffusion（history-guidance），换成 DiT + 潜空间扩散 做长视频（如微调 Wan2.1-T2V-1.3B 仅 20k 步、49 帧训练即可稳定 rollout 5× 到 217 帧），并沿用为长视频生成与规划机器人域的通用长滚动技术；该方法已成为世界模型”长时程稳定 + 可控性”方向的高频被引基石。NeurIPS 2024 收录（vol. 37, pp. 24081–24125）。

原始链接

一手源存档（sources/）

  • diffusion-forcing—github-readme — GitHub README 快照（sources/world-model/2024/diffusion-forcing--github-readme.md）
  • diffusion-forcing—project-page — 项目主页快照（sources/world-model/2024/diffusion-forcing--project-page.md）
  • arXiv 全文 PDF（2407.01392v4, arXiv 原文 PDF，不入 git）：见上方 arXiv 链接

关系图谱

反向链接