MIT 最新扩散模型课程：从微分方程探索 AIGC 的边界

摘要：3 月 4 日，MIT 计算机系的教授 Peter Holderrieth 和 Ezra Erives 在 YouTube 上传了其最新课程系列“Generative AI with Stochastic Differential Equations”，从数学

揭秘扩散模型背后的数学理念。

作者丨洪雨欣

编辑丨陈彩娴

3 月 4 日，MIT 计算机系的教授 Peter Holderrieth 和 Ezra Erives 在 YouTube 上传了其最新课程系列“Generative AI with Stochastic Differential Equations”，从数学的角度探讨了当前 AIGC 领域应用最广泛的两种算法：去噪扩散模型和流匹配。

生成模型将噪声逐步转换为数据来生成对象，这一系列的演化过程可以通过模拟常微分方程（ODE）或随机微分方程（SDE）实现。此外，Holderrieth 和 Erives 还讨论了如何利用深度神经网络大规模构建、训练和模拟这些微分方程。

该课程系列一共有 6 个章节，AI 科技评论针对原视频作了不改原意的整理如下：

1 流模型和扩散模型

先让我们定义一下“generate”。

比如我们想要生成一张狗的图片，在这一过程中，存在着一系列“可能是狗”的图片。

在机器学习中，通常将这种“可能合适”图片的多样性视为一个概率分布，我们称之为数据分布，这个分布会赋予那些看起来更像狗的图片更高的可能性。所以，图片的准确度会被其在数据分布下的可能性所取代。

因此，我们可以从数学上将生成任务表达为从数据分布中进行采样。

在机器学习中，我们需要数据来训练模型，包括互联网上公开可用的图像、YouTube视频、或者是蛋白质数据库。在生成建模中，一个数据集由来自数据分布的有限数量的样本组成，我们将其表示为集合 1 到集合 n。

在许多情况下，我们会希望基于特定条件来生成一个对象。例如，我们希望生成“一只狗在覆盖着雪的山坡上奔跑，背景是山脉”的图像。我们可以从条件数据分布中进行采样，其中 y 是一个条件变量，它可能是狗，可能是猫，也可能是风景。条件生成模型通常可以对任意变量进行条件设定，我们希望用不同的文本提示进行条件设定，因此，我们要寻求一个能够针对任意变量进行条件设定的单一模型。

所以，生成模型可以总结为：从数据分布中生成样本。

那么我们该如何获取样本呢？

假设我们可以从一个初始分布中获取样本，比如高斯分布，那么生成模型的目标就是将初始分布中采样得到的样本转换为数据分布的样本。在大多数应用中，我们都会将初始分布设定为一个简单的高斯分布。

接下来，我们将描述如何通过微分方程模拟来获得流模型和扩散模型。

首先，我们可以通过常微分方程（ODE）构建一个流模型，我们的目标是将一个初始分布转换为数据分布。常微分方程的存在性与唯一性定理是微分方程理论中的一个基本结果，所以在机器学习的实际过程中，常微分方程的解不仅存在，而且是唯一的。

在算法1中，我们总结了从流模型中进行采样的过程：

向量场是带有参数的神经网络。目前，我们将向量场看作是一个通用的神经网络，即一个带有参数的连续函数。要注意的一点是，神经网络是对向量场进行参数化，而非对流进行参数化。

接下来，我们可以通过随机微分方程（SDE），用相同的方式来构建一个生成模型。SDE是通过布朗运动构建而成的，布朗运动是一项源于对物理扩散过程研究的基本随机过程，相当于一种连续的随机游走。与处理常微分方程类似，我们使用随机初始分布来模拟随机微分方程。为了参数化这个SDE，我们可以简单地参数化其核心组成部分———向量场。

在算法2中，我们总结了从SDE中进行采样的过程：

一个扩散模型由一个神经网络和固定的扩散系数组成，其中神经网络的参数可以用于参数化一个向量场。当扩散系数为0时，扩散模型就成为了一个流模型。

2 搭建训练目标

我们将通过最小化均方误差来训练目标。首先，我们需要设置一个训练目标。像神经网络一样，训练目标本身也应该是一个向量场。除此之外，它应该做到我们期望神经网络能做到的事情：将噪声转换为数据。

构建训练目标的第一步是指定一条概率路径。直观来讲，概率路径规定了从初始噪声分布到数据分布的逐步插值过程。

条件概率路径是一组关于数据点的随机分布，是针对单个数据点的性质或行为。换句话说，条件概率路径会逐渐将单个数据点转换为初始分布。而边缘概率路径则是跨越整个数据点分布的性质或行为，你可以将概率路径看成分布空间中的一条轨迹，每个条件概率路径都会诱导出一个边缘概率路径。

去噪扩散模型所使用的是高斯概率路径，条件概率路径会对高斯分布和单个数据点对应的随机分布进行插值，边缘概率路径会对整个数据点的分布进行插值。

现在我们通过使用定义的概率路径来构建流模型的训练目标。

每个数据点的训练目标表示一个条件向量场，定义为ODE产生条件概率路径。我们利用ODE模拟概率路径，蓝色背景为数据分布，红色背景为高斯分布。上排是条件路径，下排是边缘概率路径。可以看出，条件向量场沿着条件概率路径，边缘向量场沿着边缘概率路径。

我们可以用连续方程来证明这一点。让我们考虑一个具有向量场训练目标的流模型，其中初始随机变量服从初始分布。连续性方程的公式为：在点 x处概率质量的变化等于向量场u处概率质量的流出减去流入。

我们刚刚成功地为流模型构建了一个训练目标，我们可以通过Fokker - Planck方程将连续性方程从ODE拓展到SDE。

如示意图所示，我们使用ODE来模拟概率路径。蓝色背景为数据分布，红色背景为高斯分布。上排是条件路径，下排是边缘概率路径。可以看出，SDE将初始分布的样本传输到条件路径的样本和边缘路径的样本。

对于扩散系数，我们可以构造一个遵循相同概率路径的SDE。在这个场景下，用条件概率路径和条件向量场替换边缘概率和边缘向量场，相同的结论仍然成立，所以我们可以借助条件得分函数来表示边缘得分函数，这样该随机微分方程就如所期望的那样“将噪声转换为数据”。

3 训练流模型和扩散模型

像之前的流模型一样，我们希望神经网络等于边际向量场。换句话说，我们希望减小神经网络和边际向量场之间的均方误差。

首先，抽取一个随机时间。其次，我们从数据集中抽取一个随机点，从边际概率路径中进行采样，例如添加一些噪声，并计算神经网络。最后，计算神经网络的输出与边际向量场之间的均方误差。我们将利用条件速度场的可处理性定义条件流匹配损失。在这里，我们使用条件向量场而不是边际向量场，这是因为我们用条件向量场的解析公式最小化上述损失。

一旦神经网络被训练好了，我们就能够对流模型进行模拟以此来得到采样，这套流程就是流匹配。

为了更靠近边缘得分，我们可以使用一个称之为得分网络的神经网络。我们可以设计一个得分匹配损失和一个条件得分匹配损失。理想情况下我们希望最小化得分匹配损失，由于我们不知道边缘得分，我们可以使用条件得分匹配损失来实现这一点。

训练完成后，我们可以选择任意扩散系数，然后模拟该随机微分方程以生成样本。理论上，在完美训练的情况下，每个扩散系数都应该给出服从真实数据分布的样本。但在实际中，我们会遇到两类错误：第一，由于对SDE的模拟不完善而产生的数值误差；第二，训练误差，即训练模型并不完全等同于目标模型。因此，存在一个最优但未知的噪声水平系数——这可以通过实际测试不同的值来凭经验确定。

我们可以通过高斯概率路径来训练得分匹配模型，条件得分匹配损失也被称为去噪得分匹配，它是用于学习扩散模型的最早程序之一。我们可以看到，当beta接近0时，条件得分损失在数值上是不稳定的，也就是说，只有添加足够多的噪声，去噪得分匹配才有效。

虽然流匹配仅允许通过ODE进行确定性的模拟过程，但去噪扩散模型允许进行确定性（概率流常微分方程）或者随机性（随机微分方程采样）的模拟。然而，不同于流匹配或随机插值，后者能够通过任意的概率路径将初始分布转换为数据分布。

去噪扩散模型仅适用于高斯初始分布和高斯概率路径。对于高斯概率路径来说，我们无需分别训练向量场和得分网络，它们可以在训练后相互转换。

文献中流行的扩散模型的替代公式有：

离散时间：通常使用通过离散时间得出SDE的近似值。

倒置时间约定：一种流行的做法是采用倒置时间约定，其中时间对应数据分布。

正向过程：正向过程（或加噪过程）是构建高斯概率路径的方法。

通过时间反转构建训练目标：也可以通过对SDE进行时间反转来构建训练目标。

4 搭建图像生成器

假如要构建Stable Diffusion 3和Meta Movie Gen Video这样的图像生成模型，我们首先需要构建条件生成引导机制。我们还将了解无分类器引导，这是一种用于提高条件生成质量的流行技术。

我们选择使用“引导”一词来代替“条件”，以指代基于y进行条件化的操作。引导式生成建模的目标是能够针对任意y从数据分布中进行采样。

在流匹配和得分匹配的语言体系中，我们定义一个引导扩散模型是由一个引导向量场和一个时间相关的扩散系数共同组成的。如果我们设想固定标签y的取值，那么我们就回到了无引导的生成问题。我们可以相应地使用条件流匹配目标来构建一个生成模型。

请注意，因为y不影响条件概率路径或条件向量场，所以在所有y以及所有时间里面都可以得到一个引导条件流匹配目标。引导目标与无引导目标的主要区别之一在于，引导目标是从数据分布中采样z和y，而不仅仅是采样z。原因在于，我们的数据分布现在原则上是一个关于图像和文本提示的联合分布，这样会让生成过程更准确。

虽然上述条件训练程序在理论上是有效的，但这种程序得到的图像样本与期望标签的拟合度不够好。由于人为强化引导变量 y的影响可以提高感知质量，这一见解被提炼成一种无分类器引导的技术，该技术在最先进的扩散模型领域得到了广泛应用。

为简单起见，我们在此仅关注高斯概率路径的情况。在一个模型中同时训练条件模型和无条件模型，这被称为无分类器引导。算法8展示了如何将无分类器引导构造扩展到扩散模型的环境中：

在扩散模型的开发中，我们需要介绍一种特定类型的卷积神经网络U-Net。它最初是为图像分割而设计的，其关键特征在于输入和输出都具有图像的形状，这使得它非常适合将向量场参数化。

U-Net由一系列编码器Ei、相应的一系列解码器Di以及位于它们之间的一个潜在处理块组成，我们将这个潜在处理块称为中编码器。随着输入通过编码器，其表示形式中的通道数量增加，而图像的高度和宽度减小。编码器和解码器通常都由一系列卷积层（其间包含激活函数、池化操作等）组成。输入在到达第一个编码器块之前，通常会先被送入一个初始预编码块以增加通道数量。

U-Net的一种替代方案是扩散Transformer（DiTs），它摒弃了卷积操作，纯粹使用注意力机制。扩散Transformer基于视觉Transformer（ViTs），其主要思想本质上是将图像分割成多个图像块，对每个图像块进行嵌入，然后在图像块之间进行注意力计算。例如，Stable Diffusion 3就采用条件流匹配进行训练，它将速度场参数化为一种改进的DiT。

大规模应用的一个常见问题是数据维度极高，导致消耗过多内存。例如，我们想要生成分辨率为1000×10000像素的高分辨率图像，这会产生100万个维度。

为了减少内存使用量，一种常见的设计模式是在潜在空间中进行操作，该潜在空间可被视为分辨率较低的数据的压缩版本。

具体而言，通常的方法是将流模型或扩散模型与自编码器相结合。首先通过自编码器将训练数据集编码到潜在空间中，然后在潜在空间中训练流模型或扩散模型。

采样时，首先使用训练好的流模型或扩散模型在潜在空间中进行采样，然后通过解码器对输出进行解码。

直观地说，一个训练良好的自编码器能够过滤掉语义上无意义的细节，从而使生成模型能够聚焦于重要的、感知上相关的特征。

到目前为止，几乎所有用于图像和视频生成最先进的方法都涉及在自编码器的潜在空间中训练流模型或扩散模型——这就是所谓的潜在扩散模型。

然而，需要注意的是，在训练扩散模型之前也需要训练自编码器，模型的性能也取决于自编码器将图像压缩到潜在空间以及恢复美观图像的能力。

Stable Diffusion 3 使用了我们在这项研究中的条件流匹配目标。正如他们的论文所概述的，他们对各种流和扩散替代方案进行了广泛的测试，并发现流匹配表现最佳。在训练方面，它使用了无分类器指导训练。

为了增强文本条件作用，Stable Diffusion 3 利用了三种不同类型的文本嵌入，其中包括 CLIP 嵌入和谷歌 T5-XXL 编码器预训练实例产生的序列输出。

CLIP 嵌入提供了输入文本的粗略、总体嵌入，而 T5 嵌入提供了更细粒度的上下文层次，使模型能够关注条件文本的特定元素。为了适应这些序列上下文嵌入，扩散 Transformer不仅要关注图像，还要关注文本嵌入，从而将条件能力从最初为 DiT 提出的方案扩展到序列上下文嵌入。

这种修改后的 DiT 被称为多模态 DiT（MM-DiT）。他们最大的模型拥有 80 亿个参数。

在采样方面，他们使用 50 步法则（评估网络50次），采用欧拉模拟方案，并使用 2.0-5.0 之间的无分类器指导权重。

Movie Gen Video利用具有相同CondOT路径的条件流匹配目标。与Stable Diffusion 3一样，Movie Gen Video也在冻结的预训练自编码器的潜在空间中运行。

值得关注的是，为了减少内存消耗，自编码器对视频来说比图像更加重要——这就是为什么目前大多数视频生成器在生成视频的长度方面相当受限的原因。通过引入一个时间自编码器（TAE）来处理增加的时间维度，该自编码器将原始视频映射到潜在空间。

为了适应长视频，一种时间平铺程序程序会将视频切分成片段，每个片段分别编码后拼接在一起。模型本身由一个类似DiT的主干网络给出，xt沿时间和空间维度被分块，然后将图像块传递给一个Transformer，Transformer会采用图像块之间的自注意力以及与语言模型嵌入的交叉注意力。

对于文本条件，Movie Gen Video采用了三种类型的文本嵌入：UL2嵌入，用于细粒度的基于文本的推理；ByT5嵌入，用于关注字符级细节；以及MetaCLIP嵌入，在共享的文本 - 图像嵌入空间中进行训练。他们最大的模型有300亿个参数。

5 应用于机器人的扩散模型

扩散模型会根据机器人的观察结果进行条件化，当对未来轨迹上的航点进行去噪时，这个航点上的每个点都是 10 赫兹的command，一旦完成了下一组推理，就会切换到新的组，所以这个扩散过程是循环的。当开始噪声处理时，轨迹是随机噪声，扩散模型会将其扩散成连贯的东西，你可以用这个基本配方做很多事情。

我们使用RIS相机时，有些东西无法从静态相机中轻易看到，但使用腕戴式相机时，它可以和场景相机之间共享信息，实际上腕戴式相机充当了图像随机化的功能，它总是在移动而且多样性很高，这在很大程度上解决了静态场景相机的问题。

我们通过安全层发送指令，然后传输给机器人，过程中会有一个正在填充的低级命令的缓冲区，未来的 16 或 32 个动作将会转储到缓冲区，然后立即再次开始推理。如果在推理运行之前执行了缓冲区中一半的动作，我们会覆盖缓冲区中的内容并继续。如果机器上可能还有其他程序在运行，可以在刷新缓冲区之前先执行更远的命令。

机器人会把“手移动到某个特定点”转换成一个命令，命令的力度越大，施加的力度就越大，这里的命令听起来像位置的移动，但实际上却是力度的感知。我们给它提供了一点历史记录，它可以查看之前命令了什么，在哪里命令的，两者之间的增量也可以让它推断出力度。

我们在这些触觉传感器上做了相当多的工作，我们使用了几种不同的版本，一种是用一个凝胶或充气膜，后面有一个摄像头或其他类似的传感器。例如拧紧瓶盖，实际上机器人很难感知到何时才算完全拧紧，装上这些传感器后，当感觉到瓶子变紧时，它自然就会停止。

扩散模型非常巧妙，因为它可以对多模态分布进行建模，所以在图像里，这些模型通常不会模糊模式之间的界限，你会得到一个清晰的样本。目前我们面对的问题是长视界多模态（Long Horizon Modality），这些较长的步骤可能具有不同的顺序，扩散模型在这方面表现不佳，因为它们没有足够的历史记录来了解它们应该采用哪种模式。

如果机器人操纵的物体不同，可以在训练数据中显示这种变化，这样就不会每次教一些新东西时都需要重新建立它的认知，我们在原型中也可以手动添加更多数据去修补错误，但这样做会比较麻烦。

想提高稳健性和性能，就需要使用大量数据进行多任务学习，这是一个非常有效的方法，唯一的问题机器人数据的缺乏，我们得从多样化的来源获取更多的数据。当训练数据足够大，特征空间足够丰富时，我们就能做出很多精彩的东西。

我们还做了一些流匹配的内部工作，确保匹配过程不会改变机器人的行为。

6 应用于蛋白质设计的扩散模型

我将讨论蛋白质生成的扩散模型和蛋白质设计的下游模型。

为什么我们需要AI来生成蛋白质？

打个比方，我们要设计一种叫做抗体的蛋白质，这些蛋白质会攻击入侵身体的病毒。如果我们能使用AI立即产生抗体，事情便会简单很多。数据显示，研发一种药物大约需要10年时间，约26亿美元，AI也许可以减少药物研发的时间和精力。

一个生成蛋白质的算法是从结构到序列的逆向而行，这是一种基于神经网络的模型。如果要真正生产和合成蛋白质，需要创建一个结构并筛选出一个序列，扩散模型实现了这一步。

我们希望人工智能可以快速产出新的分子库，并在实验中进行高效的微调。如果能做到这一点，我们就可以真正将生物学与人工智能联系起来，并拥有这个迭代循环的过程。一旦我们拥有结构生成的基本算法，我们就可以利用它们进行下游蛋白质的设计任务。

那么我们如何构建蛋白质结构模型呢？我们有三个选择。第一，给每个原子构建一个3D坐标，这样做的好处是我们可以精准控制每个原子的位置。坏处是，键是不固定的，导致模型难以扩展，性能不佳。

第二种选择是只构建扭转角，对角度进行扩散，这就是所谓的杠杆臂效应。这样做可以固定键，但是问题又来了，我们没办法控制原子的位置。

现在来看第三种选择，即沿着链建模框架。这样4个键当中可以固定3个，同时还能精确控制框架的位置。

我们可以用 SE（3）来表示固定键，这是一个特殊的unclean group，它既有平移又有旋转的功能。我们可以通过将两个向量从中心坐标移到其他两个坐标。这样既保持了键的固定，又可以随意移动原子的位置。

但是，当你将原子组合在一起时，你只是保存了一个常数因子，所以都会产生扩展方面的问题。蛋白质框架的扩散，就好比我们将元素组合在一起子，从噪声开始，然后是旋转和平移元素，再到噪声和数据空间之间移动。这四个过程由微分方程控制，并有反向采样的过程。

这里有一点要强调的是，我们在黎曼流形上进行扩散时必须弄清楚如何进行旋转，比如在 SE（3）上旋转。当你给蛋白质添加噪音时，会产生了一个独特的分布——SE（3）公式上的各向同性作用。当你向它添加更多的噪音时，它会开始在黎曼流形周围采样。

所谓的帧扩散就是以这种方式参数化蛋白质，我们将用噪音来破坏它们，然后训练一个模型来消除这种噪音，使其恢复到原始状态，最后从纯噪音开始，使用神经网络迭代地采样数据。

在模型架构中，我们不仅要关心标记，还要关心每个残基和每帧之间是如何互相作用的。当从扩散模型中采样时，想要保持SE（3）的不变性，每个扩散模型都得是一个等变分数模型，这样你的分数就会随着蛋白质的旋转而旋转。如果我从一些噪声里采样，最终将其解码成蛋白质，我可能会旋转这个噪音，而最终的蛋白质也可能会旋转。

从纯噪音开始生成一个随机蛋白质，模型会解码产生一组不同的蛋白质。

如何检查这些蛋白质是否良好？

我们可以通过实际性检查来看人工智能生成结构的序列。假设这个新蛋白质是由扩散模型生成的，我们将为它计算一个序列，我们会使用神经网络来预测实际结构，这是我们目前评估 AI 生成结构的主要方法。

如果不想每次都只对同一种蛋白质进行采样，有一种方法是执行聚类算法，只需取两个蛋白质计算，将其标准化为某个长度，然后执行算法。

当你得到更长的蛋白质时，性能可能会变得更差，但模型实际上能够超越训练集进行推广，通过学习等变量扩散模型，从而对训练内容非常不同的东西进行采样，这是一个非常好的迹象。

在初始化时，你必须确定模型要使用多少个token，然后给它的位置编码，比如从1 到 200 或 1 到 400 的位置编码，模型会查看位置编码并决定蛋白质生成的长度。位置编码和注意力机制也非常重要，因为token会在扩散过程中进行自我排序。

接下来是用于蛋白质生成的RF扩散模型。我们都听说过预训练对语言模型的重要性，所以我们采用预训练的蛋白质结构预测神经网络，然后用扩散对其进行微调。

预训练实际上非常关键，因为它增加了生成的长度。之前只能达到 300 的长度，通过有预训练的RF2模型可以达到800 或甚至 1000 的长度，并生成大量蛋白质。这个模型有一个0.5的临界点，如果低于这个值，那么生成的结果会与训练集中的任何东西都不太相似。最重要的一点是，RF2是针对条件生成进行训练的。

现在我们来讲对称生成，如果添加对称噪声，模型将会生成对称蛋白质。

比如，这里有一个病毒，它会生成一个与它结合的蛋白质。你也可以给它一个二维指令让其生成一个特定的蛋白质。或者可以做一个修复任务，让模型只生成一部分蛋白质。你还可以将多个条件结合在一起，添加一些额外的术语，让生成偏离它所训练的内容，转而偏向更理想的蛋白质。

但在过程中，模型可能会生成一种蛋白质，这种蛋白质可能会与小分子发生冲突。但通过对称约束，你可以让模型与分子完美结合。我们可以告诉模型，让它能够按照指令生成对称的蛋白质。

一个现有的蛋白质部分，模型会遵循对称约束来填充这个蛋白质的其余部分。重要的一点是我们必须给它提供条件，模型将遵循这些指令在特定位置结合，否则模型将结合在某个随机位置。