↓ 跳过正文

流匹配与扩散模型

·3021 字·7 分钟· loading · loading · · · 草稿
目录
AI工程 - 这篇文章属于一个选集。
§ 2: 本文

去噪扩散模型与流匹配模型是如今顶尖图像、音频和视频生成模型的核心基石。因此,熟悉相关技术并建立直观的物理与数学直觉,既十分必要,又非常有趣。

本文主要是我在学习 MIT 一门极棒的课程 MIT 6.S184 时的学习笔记。

概述
#

首先,我们的初衷是想“生成点东西”,比如一张图片或者一段视频。这是我们工作的最初动机。但究竟什么是所谓的“东西”?在数学上我们又该如何形式化地表达它?

实际上,所有我们感兴趣的对象都可以表示为一个向量 \(z \in \mathbb{R}^d\),或者能够展平为一个向量。因此,与其说“我想要一张图片”,现在我们可以说“我想要这张图片对应的向量”。这就引出了核心观点之一:

万物皆向量(An object is a vector)

下一个不太严谨的词是“生成(generation)”。直观上讲,当你想到“一张小狗的图片”时,脑海中会浮现出无数张不同小狗的模样,你很难说哪一张才是唯一正确的,只能评判哪一张更好看、更合理。注意到这与从某个数据分布中抽取一个样本的行为如出一辙,我们就可以得出结论:“生成”本质上就是“采样”。所以这里的第二个核心思想是:

生成即采样(Generation can be viewed as sampling)

现在目标就清晰多了:我们需要找到目标向量所在的概率分布。再往前推一步,当我们想要生成不同的东西时,我们可能需要“切换”分布——比如从“小狗分布”切换到“小猫分布”。

因此,最终形式化的目标,就是从条件分布 \(p_{\text{data}}(\cdot|y)\) 中采样一个向量 \(z\),其中 \(y\) 是描述我们需求的提示词(prompt)。

流模型与扩散模型
#

如果我们仅仅想生成小狗的图片,该如何表示这个分布呢?这是最终目标的一个简化版本。最朴素的想法是用一个神经网络来充当这个分布:它就像一台吃进随机种子(random seeds)就能吐出小狗图片的老虎机。

这个想法非常自然,也正是过去十年里大家一直在探索的路线。但最核心的挑战在于:如何将随机种子与真实的小狗图片绑定起来?比如摆在你面前的是一张哈士奇的照片,那它对应的随机种子到底该是哪一个?如果你忽略这一点,直接用均方误差(MSE)损失函数去暴力训练模型,你最终只会得到一只“平均化”的狗——糊成一团,甚至都很难辨认出那到底是不是一只狗。🤣

为了解决这个问题,我们需要重新定义挑战:其实我们根本没必要去较真哪个种子非得对应哪只“哈士奇”,我们真正需要做的,是训练我们的模型分布,使其在整体上尽可能接近真实的“小狗分布”。

这自然而然地将目标从“比对单个像素”转向了“度量两个概率分布之间的距离”。一旦我们不再纠结于像素级别的强制配对,而是引入巧妙的设计——比如对抗性判别器(GAN)或多步去噪(Diffusion),将生成分布作为一个整体拉向真实分布,那种模糊的“幽灵平均图”就会彻底消失,逼真、清晰的小狗便脱颖而出。

但问题随之而来:即便我们放弃了像素级的 MSE,转向分布匹配博弈(如 GAN),强求一个神经网络一步跨越——直接把一个简单的高斯噪声团变换成极其复杂、支离破碎的真实图像流形,这在数学上简直是残酷至极(mathematically brutal)。这种单步映射通常是极度不连续且病态(ill-conditioned)的。这种“一步到位”的跨越,正是早期生成模型饱受训练极度不稳定、模式崩溃(mode collapse)或严苛架构限制困扰的根源所在。

这正是 MIT 6.S 184 所带来的重大范式转变——**流与扩散:

与其尝试一次不可能的惊天一跃,何不架起一座连续渐进的桥梁?

如果从纯高斯噪声 \(z \sim \mathcal{N}(0, I)\) 一步跨到一张清晰的哈士奇图像 \(x \sim p_{\text{data}}\) 过于困难,我们完全可以把这段旅程拆解为沿时间轨迹 \(t \in [0, 1]\) 展开的平滑增量步骤。在路径上的任意一点,网络不再需要凭空捏造出一整只狗,它只需要告诉我们局部应该朝哪个方向推动样本一点点——无论是预测局部的速度向量(流匹配 Flow Matching),还是抹去一丁点微弱的噪声(扩散模型 Diffusion)。通过将混乱的全局非线性映射转化为一系列简单、局部稳定的回归任务,我们终于赢得了稳定的训练过程和令人惊艳的生成质量。

Note

这样做还附带了一个绝佳的好处:在流或扩散的过程中,我们的提示词(prompt)可以被极其自然、合理地融入引导。你可以想象一下,如果试图通过单步映射直达目标,要把提示词的语义和最终分布融合起来该有多么复杂。

从更宏观的视角来看,这种方法可以被视作是在时间维度上拓展“深度”,而非单纯在模型架构的层数维度上堆叠深度。

流模型
#

总结一下,我们目前的任务就是去描述噪声分布如何一步步平滑演化为目标数据分布的过程。

幸运的是(不得不说 🐸),常微分方程(ODE)恰好是描述这一演化过程的绝佳数学工具。为了保持本博客的完整性,在此我直接引用课程讲义(lecture notes) 中的核心表述。

在数学上,轨迹 \(X: [0, 1] \to \mathbb{R}^d, \, t \mapsto X_t\) 是由一个与时间相关的向量场 \(u_t: \mathbb{R}^d \to \mathbb{R}^d\) 驱动的,该向量场指定了在每个时间点和空间位置上的速度:

$$ \frac{d}{dt} X_t = u_t(X_t), \quad X_0 = x_0 $$

该方程随时间演化的解由流(flow)\(\psi_t: \mathbb{R}^d \to \mathbb{R}^d\) 来刻画,它追踪了满足 \(\frac{d}{dt}\psi_t(x_0) = u_t(\psi_t(x_0))\) 且 \(\psi_0(x_0) = x_0\) 的位置 \(\psi_t(x_0) = X_t\)。

构建生成式流模型的精妙之处在于:我们永远不需要神经网络去直接预测整条轨迹或流变换 \(\psi_t\) 本身。相反,我们只需要使用带权重 \(\theta\) 的网络去参数化局部的速度向量场 \(u_t^\theta(x) \approx u_t(x)\)。随机性的来源非常纯粹——完全来自于从简单的基准分布 \(p_{\text{init}}\)(例如标准高斯分布 \(\mathcal{N}(0, I_d)\))中采样的初始状态 \(X_0\)。我们最终的目标,仅仅是让终端状态 \(X_1 = \psi_1^\theta(X_0)\) 符合真实的数据分布 \(p_{\text{data}}\)。

在推理(采样)阶段,由于神经网络表示的向量场无法求得解析积分,生成过程是通过对 ODE 进行数值仿真来完成的。从 \(X_0 \sim \mathcal{N}(0, I_d)\) 开始,我们使用微小步长 \(h = 1/n\) 对时间进行离散化,并借助标准的欧拉法(Euler method)沿着向量场向前推进一步:

$$ X_{t+h} = X_t + h \cdot u_t^\theta(X_t) $$

从 \(t = 0\) 到 \(1\) 循环迭代这一过程,就能沿着学习到的速度流线,将初始高斯噪声逐步推演为最终的生成样本 \(X_1 \approx z \sim p_{\text{data}}\)。

扩散模型
#

大体上,扩散模型与流模型非常相似,只不过它的演化规则是随机的。

为了给确定性的轨迹注入随机性,我们引入了由标准布朗运动(或维纳过程)\(W_t\) 驱动的连续随机游走,它具有连续的轨迹和独立的高斯增量 \(W_{t+h} - W_t \sim \mathcal{N}(0, h I_d)\)。在 ODE 的无限小微元步长中加入这些随机的“微扰”,便得到了标准的随机微分方程(SDE):

$$ dX_t = u_t(X_t)dt + \sigma_t dW_t, \quad X_0 \sim p_{\text{init}} $$

在此,\(u_t(x)\) 充当确定性的漂移向量场,而 \(\sigma_t \ge 0\) 是一个标量扩散系数,负责随时间调节注入噪声的强度。

与流模型如出一辙,神经网络仅需学习漂移场 \(u_t^\theta(x)\),而 \(\sigma_t\) 通常作为预先设定好的固定调度策略。从扩散模型中采样的逻辑也与流模型高度镜像,只是将欧拉积分器替换为其随机版本——欧拉-丸山法:

$$ X_{t+h} = X_t + h \cdot u_t^\theta(X_t) + \sigma_t \sqrt{h} \cdot \epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, I_d) $$

在每个微小区间内,状态既沿着向量场方向迈出一步,同时还会受到一个由 \(\sigma_t \sqrt{h}\) 缩放的高斯随机扰动。如果我们通过令 \(\sigma_t = 0\) 完全关闭这种噪声注入,随机项便瞬间消失,我们立刻就能退化还原到确定性的流模型。这也揭示了一个重要事实:流模型本质上只是更广义 SDE 家族中扩散系数为零的一个特例。

思考

为什么我们要引入随机因素?目前一个较为公认的解释是为了多样性。流模型简单优雅,但往往过于僵硬且具有强确定性。不过,我对这种观点持有一点保留意见,这值得更深入地探讨。🤔

AI工程 - 这篇文章属于一个选集。
§ 2: 本文

相关文章