[论文解读] Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

Boyuan Chen, Diego Martí Monsó+4 more|arXiv (Cornell University)|Jul 1, 2024

Forensic and Genetic ResearchBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

该论文提出了一种名为 Diffusion Forcing 的新型训练范式，通过为每个 token 分配独立的噪声水平，将 next-token 预测与全序列扩散统一起来，实现了稳定、可变长度的生成以及长时序序列建模中的有效引导。该方法通过结合因果建模与灵活的、逐 token 的去噪调度，在视频生成、规划和时间序列预测任务中实现了最先进性能。

ABSTRACT

This paper presents Diffusion Forcing, a new training paradigm where a diffusion model is trained to denoise a set of tokens with independent per-token noise levels. We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens without fully diffusing past ones. Our approach is shown to combine the strengths of next-token prediction models, such as variable-length generation, with the strengths of full-sequence diffusion models, such as the ability to guide sampling to desirable trajectories. Our method offers a range of additional capabilities, such as (1) rolling-out sequences of continuous tokens, such as video, with lengths past the training horizon, where baselines diverge and (2) new sampling and guiding schemes that uniquely profit from Diffusion Forcing's variable-horizon and causal architecture, and which lead to marked performance gains in decision-making and planning tasks. In addition to its empirical success, our method is proven to optimize a variational lower bound on the likelihoods of all subsequences of tokens drawn from the true joint distribution. Project website: https://boyuan.space/diffusion-forcing

研究动机与目标

• 为解决自回归模型在生成连续数据（如视频）的长序列时，超出训练时序范围后出现的不稳定性问题。
• 克服全序列扩散模型缺乏因果性、无法支持可变长度生成或子序列条件生成的局限性。
• 通过结合 next-token 预测的灵活性与扩散模型的全局优化能力，实现有效的序列引导与规划。
• 形式化一个训练目标，以最大化真实联合分布中所有子序列的似然性的变分下界。
• 在包括视频生成、基于模型的规划和时间序列预测在内的多样化领域中，验证该方法的有效性。

提出的方法

• Diffusion Forcing 训练单一因果 next-token 预测模型，使其对每个 token 独立施加特定噪声水平的噪声，从而实现逐 token 的去噪调度。
• 该模型被训练为从带噪声的序列中重建任意 token 子集，从而有效学习根据每个 token 的独立噪声水平实现‘去掩码’。
• 推理阶段，模型对完整序列执行渐进式去噪，噪声水平可任意设定为每个 token 独立的值，从而实现从下一个 token 开始的可变长度生成。
• 提出一种新型采样方案——蒙特卡洛树引导（Monte Carlo Tree Guidance, MCTG），利用因果架构与可变时序能力，在规划任务中引导采样向高回报轨迹前进。
• 该方法采用共享的因果 Transformer 架构，使未来 token 的预测基于历史 token，既保持因果性，又支持全序列去噪。
• 训练目标优化了在训练数据中观察到的所有子序列的似然性的变分下界，确保与联合分布的一致性。

Figure 2 : Method Overview. Diffusion Forcing trains causal sequence neural networks (such as an RNN or a masked transformer) to denoise flexible-length sequences where each frame of the sequence can have a different noise level. In contrast, next-token prediction models, common in language modeling

实验结果

研究问题

• RQ1 一个单一模型能否结合 next-token 预测的可变长度生成能力与全序列扩散的长时序引导能力？
• RQ2 逐 token 的噪声调度是否能实现对视频等连续序列的稳定生成，且在训练时序范围之外仍保持稳定？
• RQ3 能否在扩散框架中利用 next-token 模型的因果架构，以支持新型引导方案（如蒙特卡洛树引导）？
• RQ4 所提出的训练目标是否在理论上通过优化子序列似然性的下界而得到合理化？
• RQ5 Diffusion Forcing 在包括视频、规划和时间序列预测在内的多样化序列建模任务中表现如何？

主要发现

• Diffusion Forcing 实现了对连续序列（如视频）的稳定、长时序生成，可生成超过数千个 token 的序列，远超训练时序范围，而自回归基线模型则出现发散。
• 该方法在视频生成和时间序列预测任务中达到最先进性能，其 CRPS 分数在多个 GluonTS 数据集上显著低于基线模型。
• 使用 Diffusion Forcing 的蒙特卡洛树引导（MCTG）在决策与规划任务中带来显著改进，优于非因果的全序列扩散基线。
• 该模型在新轨迹上表现出组合性泛化能力，支持高效的树搜索与在线反馈控制，融合了自回归模型与扩散模型的优势。
• 实证结果表明，训练目标能有效最大化子序列似然性的变分下界，验证了该方法的理论基础。
• 在视觉模仿学习与基于模型的规划中，Diffusion Forcing 实现了高回报轨迹的生成，同时提升了样本效率与稳定性。