机器之心原创

2024/07/23 11:04

无限生成视频，还能规划决策，扩散强制整合下一token预测与全序列扩散

当前，采用下一 token 预测范式的自回归大型语言模型已经风靡全球，同时互联网上的大量合成图像和视频也早已让我们见识到了扩散模型的强大之处。

近日，MIT CSAIL 的一个研究团队（一作为 MIT 在读博士陈博远）成功地将全序列扩散模型与下一 token 模型的强大能力统合到了一起，提出了一种训练和采样范式：Diffusion Forcing（DF）。

论文标题：Diffusion Forcing:Next-token Prediction Meets Full-Sequence Diffusion
论文地址：https://arxiv.org/pdf/2407.01392
项目网站：https://boyuan.space/diffusion-forcing
代码地址：https://github.com/buoyancy99/diffusion-forcing

如下所示，扩散强制在一致性和稳定性方面都明显胜过全序列扩散和教师强制这两种方法。

在该框架中，每个 token 都关联了一个随机的、独立的噪声水平，并且可使用一种共享的下一 token 预测模型或下几 token 预测模型根据任意的、独立的、每 token 的方案对 token 进行去噪。

该方法的研究灵感来自这一观察：对 token 加噪声的过程就是一种形式的部分掩码过程 —— 零噪声就意味着未对 token 加掩码，而完整噪声则是完全掩蔽 token。因此，DF 可强迫模型学习去除任何可变有噪声 token 集合的掩码（图 2）。

与此同时，通过将预测方法参数化为多个下一 token 预测模型的组合，该系统可以灵活地生成不同长度的序列，并以组合方式泛化到新的轨迹（图 1）。

该团队将用于序列生成的 DF 实现成了因果扩散强制（Causal Diffusion Forcing/CDF），其中未来 token 通过一个因果架构依赖于过去 token。他们训练该模型一次性去噪序列的所有 token（其中每个 token 都有独立的噪声水平）。

在采样期间，CDF 会将一个高斯噪声帧序列逐渐地去噪成洁净的样本，其中不同帧在每个去噪步骤可能会有不同的噪声水平。类似于下一 token 预测模型，CDF 可以生成长度可变的序列；不同于下一 token 预测，CDF 的表现非常稳定 —— 不管是预测接下来的一个 token，还是未来的数千 token，甚至是连续 token。

此外，类似于全序列扩散，它也可接收引导，从而实现高奖励生成。通过协同利用因果关系、灵活的范围和可变噪声调度，CDF 能实现一项新功能：蒙特卡洛树引导（MCTG）。相比于非因果全序列扩散模型，MCTG 能极大提升高奖励生成的采样率。图 1 给出了这些能力的概况。

Diffusion Forcing（扩散强制）

1、将加噪过程视为部分掩码

首先，我们可以将任意 token 集合（不管是否为序列）视为一个通过 t 索引的有序集合。那么，使用教师强制（teacher forcing）训练下一 token 预测便可被解释成掩蔽掉时间 t 的每个 token x_t 基于过去 x_{1:t−1} 预测它们。

对于序列，可将这种操作描述成：沿时间轴执行掩码。我们可以将全序列前向扩散（即逐渐向数据

添加噪声的过程）看作一种部分掩码（partial masking），这可被称为「沿噪声轴执行掩码）。

事实上，在 K 步加噪之后，

（大概）就是白噪声了，不再有任何有关原数据的信息。

如图 2 所示，该团队建立了一个统一视角来看待沿这两个轴的掩码。

2、扩散强制：不同 token 的噪声水平不同

扩散强制（DF）框架可用于训练和采样任意序列长度的有噪声 token

，其中的关键在于每个 token 的噪声水平 k_t 会随时间步骤而变化。

这篇论文关注的重点是时间序列数据，因此他们通过一种因果架构实例化了 DF，并由此得到了因果扩散强制（CDF）。简单来说，这是使用基础循环神经网络（RNN）获得的一种最小实现。

权重为 θ 的 RNN 维护着获悉过去 token 影响的隐藏状态 z_t，其会通过一个循环层根据动态

而演化。当获得输入噪声观察

时，就以马尔可夫方式更新该隐藏状态。

当 k_t=0 时，这就是贝叶斯过滤中的后验更新；而当 k_t= K（纯噪声、无信息）时，这就等价于建模贝叶斯过滤中的「后验分布」p_θ(z_t | z_{t−1})。

给定隐藏状态 z_t，观察模型 p_θ(x_t^0 | z_t) 的目标是预测 x_t；这个单元的输入 - 输出行为与标准的条件扩散模型一样：以条件变量 z_{t−1} 和有噪声 token 为输入，预测无噪声的 x_t=x_t^0，并由此间接地通过仿射重新参数化预测噪声 ε^{k_t}。因此，我们就可以直接使用经典的扩散目标来训练（因果）扩散强制。根据噪声预测结果 ε_θ，可以对上述单元进行参数化。然后，通过最小化以下损失来找到参数 θ：

算法 1 给出了伪代码。重点在于，该损失捕获了贝叶斯过滤和条件扩散的关键元素。该团队也进一步重新推断了用于扩散强制的扩散模型训练中的常用技术，详见原论文的附录部分。他们也得出了一个非正式的定理。

定理 3.1（非正式）。扩散强制训练流程（算法 1）是在期望对数似然

上优化证据下限（ELBO）的重新加权，其中期望值会在噪声水平上平均，而

是根据前向过程加噪。此外，在适当条件下，优化 (3.1) 式还可以同时最大化所有噪声水平序列的似然下限。

扩散强制采样和所得到的能力

算法 2 描述了采样过程，其定义是：在二维的 M × T 网格 K ∈ [K]^{M×T} 上指定噪声调度；其中列对应于时间步骤 t，m 索引的行则决定了噪声水平。

为了生成长度为 T 的整个序列，先将 token x_{1:T} 初始化为白噪声，对应于噪声水平 k = K。然后沿着网格逐行向下迭代，并从左到右逐列去噪，直到噪声水平达到 K。到最后一行 m = 0 时，token 的噪声已清理干净，即噪声水平为 K_{0,t} ≡ 0。

这个采样范式会带来如下新能力：

让自回归生成变得稳定
保持未来的不确定
长期引导能力

将扩散强制用于灵活的序列决策

扩散强制的新能力也带来了新的可能性。该团队基于此为序列决策（SDM）设计了一种全新框架，并且将其成功应用到了机器人和自主智能体领域。

首先，定义一个马尔可夫决策过程，该过程具有动态 p (s_{t+1}|s_t, a_t)、观察 p (o_t|s_t) 和奖励 p (r_t|s_t, a_t)。这里的目标是训练一个策略 π(a_t|o_{1:t})，使得轨迹

的预期累积奖励最大化。这里分配 token x_t = [a_t, r_t, o_{t+1}]。一条轨迹就是一个序列 x_{1:T}，其长度可能是可变的；训练方式则如算法 1 所示。

在执行过程的每一步 t，都有一个隐藏状态 z_{t-1} 总结过去的无噪声 token x_{1:t-1}。基于这个隐藏状态，根据算法 2 采样一个规划

，其中

包含预测的动作、奖励和观察。H 是一个前向观察窗口，类似于模型预测控制中的未来预测。在采用了规划的动作之后，环境会得到一个奖励和下一个观察，从而得到下一个 token。其中隐藏状态可以根据后验 p_θ(z_t|z_{t−1}, x_t, 0) 获得更新。

该框架既可以作为策略，也可以作为规划器，其优势包括：

具有灵活的规划范围
可实现灵活的奖励引导
能实现蒙特卡洛树引导（MCTG），从而实现未来不确定性

实验

该团队评估了扩散强制作为生成序列模型的优势，其中涉及视频和时间序列预测、规划和模仿学习等多种应用。

视频预测：一致且稳定的序列生成和无限展开

针对视频生成式建模任务，他们基于 Minecraft 游戏视频和 DMLab 导航为因果扩散强制训练了一个卷积 RNN 实现。

图 3 展示了扩散强制与基准的定性结果。

可以看到，扩散强制能稳定地展开，甚至能超过其训练范围；而教师强制和全序列扩散基准会很快发散。

扩散规划：MCTG、因果不确定性、灵活的范围控制

扩散强制的能力能为决策带来独有的好处。该团队使用一种标准的离线强化学习框架 D4RL 评估了新提出的决策框架。

表 1 给出了定性和定量的评估结果。可以看到，扩散强制在全部 6 个环境中都优于 Diffuser 和所有基准。

可控的序列组合生成

该团队发现，仅需修改采样方案，就可以灵活地组合训练时间观察到的序列的子序列。

他们使用一个 2D 轨迹数据集进行了实验：在一个方形平面上，所有轨迹都是始于一角并最终到达对角，形成一种十字形。

如上图 1 所示，当不需要组合行为时，可让 DF 保持完整记忆，复制十字形的分布。当需要组合时，可让模型使用 MPC 无记忆地生成更短的规划，从而实现对这个十字形的子轨迹的缝合，得到 V 形轨迹。

机器人：长范围模仿学习和稳健的视觉运动控制

扩散强制也为真实机器人的视觉运动控制带来了新的机会。

模仿学习是一种常用的机器人操控技术，即学习专家演示的观察到动作的映射。但是，缺乏记忆往往会让模仿学习难以完成长范围的任务。DF 不仅能缓解这个短板，还能让模仿学习更稳健。

使用记忆进行模仿学习。通过遥控 Franka 机器人，该团队收集了一个视频和动作数据集。如图 4 所示，任务就是利用第三个位置交换苹果和橘子的位置。水果的初始位置是随机的，因此可能的目标状态有两个。

此外，当第三个位置有一个水果时，就无法通过当前观察推断出所需结果 —— 策略必须记住初始配置才能决定移动哪个水果。不同于常用的行为克隆方法，DF 可以自然地将记忆整合进自己的隐藏状态中。结果发现，DF 能实现 80% 的成功率，而扩散策略（当前最佳的无记忆模仿学习算法）却失败了。

此外，DF 还能更稳健地应对噪声并助益机器人预训练。

时间序列预测：扩散强制是一种优秀的通用序列模型

对于多变量时间序列预测任务，该团队的研究表明 DF 足以与之前的扩散模型和基于 Transformer 的模型媲美。

更多技术细节和实验结果请参阅原论文。

工程Diffusion ForcingMIT CSAIL

相关数据

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源，如线程、进程或数据流；也可以指硬件资源，如处理器、网络连接或扩展卡。进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态，允许多位用户有效地同时共享系统资源，或达到指定的服务质量。 see planning for more details

来源：维基百科

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

规划技术

人工智能领域的「规划」通常是指智能体执行的任务/动作的自动规划和调度，其目的是进行资源的优化。常见的规划方法包括经典规划（Classical Planning）、分层任务网络（HTN）和 logistics 规划。

来源：机器之心

伪代码技术

伪代码，又称为虚拟代码，是高层次描述算法的一种方法。它不是一种现实存在的编程语言；它可能综合使用多种编程语言的语法、保留字，甚至会用到自然语言。它以编程语言的书写形式指明算法的职能。相比于程序语言它更类似自然语言。它是半形式化、不标准的语言。

来源：维基百科

映射技术

映射指的是具有某种特殊结构的函数，或泛指类函数思想的范畴论中的态射。逻辑和图论中也有一些不太常规的用法。其数学定义为：两个非空集合A与B间存在着对应关系f，而且对于A中的每一个元素x，B中总有有唯一的一个元素y与它对应，就这种对应为从A到B的映射，记作f：A→B。其中，y称为元素x在映射f下的象，记作：y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域，记作f(A)。同样的，在机器学习中，映射就是输入与输出之间的对应关系。

来源：Wikipedia

马尔可夫决策过程技术

马尔可夫决策过程为决策者在随机环境下做出决策提供了数学架构模型，为动态规划与强化学习的最优化问题提供了有效的数学工具，广泛用于机器人学、自动化控制、经济学、以及工业界等领域。当我们提及马尔可夫决策过程时，我们一般特指其在离散时间中的随机控制过程：即对于每个时间节点，当该过程处于某状态(s)时，决策者可采取在该状态下被允许的任意决策(a)，此后下一步系统状态将随机产生，同时回馈给决策者相应的期望值，该状态转移具有马尔可夫性质。

来源：Reinforcement Learning: An Introduction, by R. Sutton and A. G. Barto； Algorithms for Reinforcement Learning, by C. Szepesvari

模仿学习技术

模仿学习（Imitation Learning）背后的原理是是通过隐含地给学习器关于这个世界的先验信息，就能执行、学习人类行为。在模仿学习任务中，智能体（agent）为了学习到策略从而尽可能像人类专家那样执行一种行为，它会寻找一种最佳的方式来使用由该专家示范的训练集（输入-输出对）。

来源：机器之心

强化学习技术

强化学习是一种试错方法，其目标是让软件智能体在特定环境中能够采取回报最大化的行为。强化学习在马尔可夫决策过程环境中主要使用的技术是动态规划（Dynamic Programming）。流行的强化学习方法包括自适应动态规划（ADP）、时间差分（TD）学习、状态-动作-回报-状态-动作（SARSA）算法、Q 学习、深度强化学习（DQN）；其应用包括下棋类游戏、机器人控制和工作调度等。

来源：机器之心

时间序列预测技术

时间序列预测法其实是一种回归预测方法，属于定量预测，其基本原理是;一方面承认事物发展的延续性，运用过去时间序列的数据进行统计分析，推测出事物的发展趋势；另一方面充分考虑到偶然因素影响而产生的随机性，为了消除随机波动的影响，利用历史数据进行统计分析，并对数据进行适当处理，进行趋势预测。

来源：百度百科

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。由于字词与句子都是任意组合的长度，因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题)，也使得在语料库中估算字串的机率变得很困难，这也是要使用近似的平滑n元语法(N-gram)模型之原因。

来源：维基百科

视频生成技术

视频生成是指利用深度学习等技术生成视频的任务。

视频预测技术

视频预测任务是指在基于给定的视频帧序列，生成一个可信的未来帧序列。