MIT 6.S184 Introduction to Flow Matching and Diffusion Models 学习笔记

前言

受 CSDIY 推荐，觉得还不错的课程。25 年的 B 站资源链接：BV1gc8Ez8EFL。

虽然我看的是 26 年版本的 lecture notes。

从微分方程的视角讲解了 flow 和 diffusion。

基本概念

需要生成的对象（图片，视频，蛋白质）等都可以使用向量 $z\in\mathbb{R}^{d}$ 来表示。

生成即采样：“生成”一个对象本质上就是在数据分布 $p_{\text{data}}$ 里面采样。

数据集：$z_1,z_2,\cdots,z_n\sim p_{\text{data}}$。

条件生成：sample 的时候带上某种条件，即 $z\sim p_{\text{data}}(\cdot | y)$，其中 $y$ 为控制条件的变量。

我们的目标就是，通过一个模型，将从已知的初始分布 $p_{\text{init}}$ 里面采样的 $x$ 转换成 $p_{\text{data}}$ 里的采样。

Flow 和 Diffusion

这一节介绍了 ODE 和 SDE 两个数学工具。

Flow 与 ODE

首先定义 ODE 的解，即 trajectory：

$$X:[0,1]\to \mathbb{R}^d,\quad t \mapsto X_t$$

即将 $[0,1]$ 中的时间 $t$ 映射到 $\mathbb{R}^d$ 空间中的某个位置。

ODE 是由向量场定义的：

$$u: \mathbb{R}^d\times [0,1] \to \mathbb{R}^d,\quad (x,t)\mapsto u_t(x)$$

相当于在空间中每个位置的每个时间点都有定义一个类似于速度的东西。ODE 就是希望我们的轨迹是沿着这个向量场走的：

$$\begin{aligned}\frac{\mathrm{d}}{\mathrm{d} t}X_t &= u_t(X_t)\\ X_0 &= x_0 \end{aligned}$$

上面的式子是相当于，$X_t$ 关于时间的导数是由在此位置的向量场给出的；下面的式子则是规定了一个初始条件。

接下来定义 flow：flow 即是把所有初值的解“打包”成一个映射：

$$\psi:\mathbb{R}^d\times[0,1]\to\mathbb{R}^d,\quad (x_0,t)\mapsto \psi_t(x_0)\\ \begin{aligned}\frac{\mathrm{d}}{\mathrm{d} t}\psi_t(x_0) &= u_t(\psi_t(x_0))\\ \psi_0(x_0) &= x_0 \end{aligned}$$

这告诉我们，对于初始位置 $X_0=x_0$，一个 ODE 的解 $X_t$ 可以由流 $\psi_t(x_0)$ 给出。

流的存在性与唯一性定理：若向量场 $u$ 连续可微且导数有界，则 ODE 有解且唯一（即 $\psi_t$），这种情况下 $\psi_t$ 是微分同胚。

这个定理的直觉告诉我们刚才给出的流是良定义的，一个初始点对应一条轨迹，且轨迹之间是不会相交的，并且这个轨迹光滑且可逆。

在机器学习中，这个 $u$ 的保证几乎一定成立（通常我们用神经网络建模 $u$）。

一般而言，求解 ODE 是没那么简单的，不过可以进行数值上的模拟。最简单的欧拉法：

$$X_{t+h}=X_t+hu_t(X_t)\quad(t = 0,h,2h,\cdots, 1-h)$$

有了上面这些工具，我们回忆一下初始目标：把 $p_{\text{init}}$ 转换成 $p_{\text{data}}$。很自然的想法就是利用 ODE 和 flow，而一个 flow model 如下被定义：

$$X_0\sim p_{\text{init}}\\ \frac{\mathrm{d}}{\mathrm{d} t}X_t = u_t^{\theta}(X_t)$$

这个向量场由参数为 $\theta$ 的神经网络定义（具体架构之后讨论）。我们的目标就是让 $X_1$ 服从数据分布：

$$X_1\sim p_{\text{data}}\iff \psi_1^{\theta}(X_0) \sim p_{\text{data}}$$

注意：神经网络建模的是向量场而不是 flow 本身。

Diffusion 与 SDE

简单来说，就是把 ODE 换成 SDE。这里先给出 SDE 的简要概念。莫名 callback 暑假的量化金融专题。？

SDE 相当于把 ODE 的确定轨迹变成了一个随机过程 $(X_t)_{0\le t\le 1}$。

在给出 SDE 之前先给出 布朗运动 的概念。布朗运动是一个随机过程 $(W_t)_{0\le t\le 1}$，其中 $W_0 = 0$，然后满足两个性质：

• Normal Increments：$W_t-W_s\sim \mathcal N(0,(t-s)I_d)$
• Independent Increments：对于任意 $0\le t_0<t_1<\cdots<t_n=1$ 都有 $W_{t_1}-W_{t_0}, W_{t_2}-W_{t_1},\cdots, W_{t_n}-W_{t_{n-1}}$ 相互独立。

SDE 的想法就是在 ODE 的基础上加上随机的布朗运动。但是布朗运动几乎处处不可导所以不太可能有上面那种利用导数给出的微分方程。我们先利用导数的语言，把 ODE 转写一下：

$$\begin{aligned}\frac{\mathrm{d}}{\mathrm{d} t}X_t &= u_t(X_t)\\ \frac 1h (X_{t+h}-X_t) &= u_t(X_t)+R_t(h)\\ X_{t+h}&= X_t + hu_t(X_t)+hR_t(h) \end{aligned}$$

其中 $R_t(h)$ 是误差余项满足 $\lim_{h\to 0}R_t(h) = 0$。我们把最后那个式子修改成随机的，即在每个微小时间步 $h$ 里面加入布朗运动的贡献：

$$X_{t+h}=X_t+hu_t(X_t)+\sigma_t(W_{t+h}-W_t)+hR_t(h)$$

其中 $\sigma_t\ge 0$ 为扩散系数，$R_t(h)$ 为随机误差项。上面这个式子就构成一个随机微分方程，可以用下面的符号来表示：

$$\begin{aligned} \mathrm{d} X_t &= u_t(X_t)\mathrm{d} t + \sigma_t \mathrm{d} W_t\\ X_0&= x_0 \end{aligned}$$

类似地，SDE 的解也有唯一性之类的性质，但此处不会进行说明。

对于 SDE 的模拟，可以使用 Euler-Maruyama 方法，和 ODE 的 Euler 方法很像：

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

通过 SDE，我们也可以构造一个生成式模型，称为扩散模型：

$$\begin{aligned} X_0&\sim p_{\text{init}} \\ \mathrm{d} X_t &= u_t(X_t)\mathrm{d} t + \sigma_t \mathrm{d} W_t \end{aligned}$$

其中 $\sigma_t: [0,1]\to [0,+\infty),t\mapsto \sigma_t$ 是一开始固定的超参。

Flow Matching

上面已经讲了 flow model 的基本概念，但没有说怎么训练这样的一个模型。接下来介绍 flow matching（流匹配）。

不过现在我们只知道 $X_0\sim p_{\text{init}}$ 以及我们希望让 $X_1\sim p_{\text{data}}$，让模型凭空学出来这个转化还是比较难的。所以 flow matching 的思路就是不直接让模型“凭空学会”如何从噪声到数据，而是先规定中间每个时刻 $t$ 的“应该长什么样”（分布路径 $p_t$），再去学习一个向量场让轨迹的分布匹配这条路径。

这一节会反复强调 conditional（条件）和 marginal（边缘）两个概念。提到条件的时候总是相当于在 $p_{\text{init}}$ 上已经取了一个 $z$ 作为条件的情况，而边缘的意思就是相当于对于所有的 $p_{\text{init}}$ 上的 $z$。

概率路径

对每个数据点 $z\in\mathbb{R}^d$，定义一条条件插值路径 $p_t(\cdot| z)$，满足端点条件：

• $t=0$ 时是噪声分布 $p_{\text{init}}$
• $t=1$ 时退化成把概率质量“钉死在 $z$”上的 Dirac $\delta_z$

原文定义为：

$$p_0(\cdot| z)=p_{\text{init}},\qquad p_1(\cdot| z)=\delta_z$$

直觉：这是“把初始噪声逐渐变成一个特定样本 $z$”的分布轨迹。

条件路径会诱导一条边缘路径 $p_t(x)$：先采样 $z\sim p_{\text{data}}$，再采样 $x\sim p_t(\cdot| z)$。这给出了一个可采样的 $p_t$：

• 采样过程：$\;z\sim p_{\text{data}},\,x\sim p_t(\cdot| z)\Rightarrow x\sim p_t$
• 概率密度：

  $$p_t(x)=\int p_t(x| z)p_{\text{data}}(z)\mathrm{d}z$$

并且由端点条件推出边缘路径确实从噪声插到数据：

$$p_0=p_{\text{init}},\qquad p_1=p_{\text{data}}$$

这告诉我们可以直接采样 $x\sim p_t$，但是无法显式计算概率密度 $p_t(x)$。

条件向量场与边缘向量场

说完了我们想要的概率路径，但还没说怎么构造出能让 $X_t$ 沿着概率路径走的向量场。本节就进行构造。

对于每个数据 $z$，定义条件向量场 $u_t^{\text{target}}(\cdot|z)$​：

$$X_0\sim p_{\text{init}},\quad \frac{\mathrm{d} X_t}{\mathrm{d} t}=u_t^{\text{target}}(X_t | z)\quad\Rightarrow\quad X_t\sim p_t(\cdot| z)$$

直觉：给定终点 $z$，这个速度场会把噪声“推向 $z$”并在每个时刻呈现出我们指定的条件分布。而且好消息是这个条件向量场一般是容易有解析解的。但是条件向量场没法给我们直接用，他只是把 $p_{\text{init}}$ 坍缩成 $\delta_z$。

但是，利用 marginalization trick，我们可以把条件向量场给利用起来，构造出边缘向量场：

定理（Marginalization trick）：

令 $u_t^{\text{target}}(x|z)$ 为条件向量场，则边缘向量场

$$u_t^{\text{target}}(x) = \int u_t^{\text{target}} (x|z) \frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}\mathrm{d}z$$

满足

$$X_0\sim p_{\text{init}},\quad \frac{\mathrm{d}}{\mathrm{d} t}X_t = u_t^{\text{target}}(X_t)\quad \Rightarrow \quad X_t\sim p_t\quad(0\le t\le 1)$$

即这样定义的边缘向量场会遵循边缘概率路径。

注意到，$u_t^{\text{target}}(x|z)$ 乘上的系数 $\displaystyle \frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}$ 实际上就是给定噪声数据 $x$，$z$ 的后验 $p_t(z|x)$。

这告诉我们 $X_1\sim p_{\text{data}}$，即这个边缘向量场可以将初始分布 $p_{\text{init}}$ 转变为数据分布 $p_{\text{data}}$。

为什么这个式子成立呢？先给出一个引理（此处不打算证明）：

连续性方程

对于向量场 $u_t^{\text{target}}$ 和 $X_0\sim p_{\text{init}}$。对于 $\forall t\in [0,1]$，$X_t\sim p_t$ 成立当且仅当

$$\partial_t p_t(x) = -\mathrm{div}(p_t u_t^{\text{target}})(x)\qquad\forall x\in\R^d,t\in[0,1]$$

其中散度 $\mathrm{div}$ 定义如下：

$$\mathrm{div}(v_t)(x) = \sum_{i=1}^d\frac{\partial}{\partial x_i}v_t^i(x)$$

直觉：

• 左边 $\partial_t p_t(x)$：某点处概率密度随时间的变化
• 右边 $-\mathrm{div}(p_t u_t)$：概率质量“流入/流出”的净效应（div 表示净流出，所以加负号是净流入）

对于边缘向量场的证明：

$$\begin{aligned} \partial_t p_t(x)&\overset{(i)}{=}\partial_t \int p_t(x | z)p_{\mathrm{data}}(z)\mathrm{d} z \\ &= \int \partial_t p_t(x | z)p_{\mathrm{data}}(z)\mathrm{d} z \\ &\overset{(ii)}{=} \int -\operatorname{div}\!\left(p_t(\cdot | z)u_t^{\mathrm{target}}(\cdot | z)\right)(x)p_{\mathrm{data}}(z)\mathrm{d} z \\ &\overset{(iii)}{=} -\operatorname{div}\!\left(\int p_t(x | z)u_t^{\mathrm{target}}(x | z)p_{\mathrm{data}}(z)\mathrm{d} z\right) \\ &\overset{(iv)}{=} -\operatorname{div}\!\left( p_t(x)\int u_t^{\mathrm{target}}(x | z) \frac{p_t(x | z)p_{\mathrm{data}}(z)}{p_t(x)}\mathrm{d} z \right)(x) \\ &\overset{(v)}{=} -\operatorname{div}\!\left(p_tu_t^{\mathrm{target}}\right)(x) \end{aligned}$$

$(i)$ 比较显然；$(ii)$ 是代入连续性方程；$(iii)$ 是根据散度的定义式交换了散度和积分的顺序；$(iv)$ 是同时乘/除 $p_t(x)$；$(v)$ 是代入 $u_t^{\text{target}}$ 的定义。

训练方法

现在我们已经知道要用神经网络学习的向量场长什么样了（即我们已经明确了训练目标）。那么现在怎么学习一个这样的向量场呢？首先需要定义误差函数。定义 $\text{Unif} = \text{Uniform}_{[0,1]}$ 为 $[0,1]$ 上的均匀分布，$\mathbb{E}$ 为期望。直观上我们希望 $u^\theta_t$ 接近 $u_t^{\text{target}}$ 所以可以用一个均方误差，定义 flow matching loss 如下：

$$\begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t\sim \text{Unif},x\sim p_t}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x) \|^2\right]\\ &\overset{(i)}{=} \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x) \|^2\right] \end{aligned}$$

其中 $p_t(x) = \int p_t(x|z)p_{\text{data}}(z)$ 是边缘概率路径，$(i)$ 则使用了 sample 的技巧。但是问题在于，我们不知道 $u_t^{\text{target}}$ 怎么算（也算不出来）。所以先定义 conditional flow matching loss 如下：

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x|z) \|^2\right]$$

注意到这个 $u_t^{\text{target}}(x|z)$ 通常是有解析解的，那么这个 loss 就好算了。但是我们要学的是边缘向量场 $u_t^{\text{target}}$ 啊？

定理：边缘流损失函数等于条件流损失函数加上与 $\theta$ 无关的常数。

$$\mathcal{L}_{\text{FM}}(\theta) = \mathcal{L}_{\text{CFM}}(\theta) + C$$

这说明他们的梯度是相同的：

$$\nabla_{\theta}\mathcal{L}_{\text{FM}}(\theta) = \nabla_{\theta}\mathcal{L}_{\text{CFM}}(\theta)$$

所以利用 SGD 优化 CFM loss 是等价于优化 FM loss 的！这说明我们只要对着 CFM 优化就可以了。

证明：利用 $\|a-b\|^2=\|a\|^2-2a^\top b + \|b\|^2$，把这两个 loss 拆开。

$$\begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\theta}(x)-u_t^{\text{target}}(x) \|^2]\\ &= \mathbb{E}_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\theta}(x)\|^2 - 2u_t^{\theta}(x)^\top u_t^{\text{target}}(x) + \|u_t^{\text{target}}(x) \|^2]\\ &= \mathbb{E}_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\theta}(x)\|^2] - 2\mathbb{E}_{t\sim\text{Unif},x\sim p_t}[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)] + \mathbb{E}_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\text{target}}(x) \|^2]\\ &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{init}},x\sim p_t(\cdot|z)}[\|u_t^{\theta}(x)\|^2] - 2\mathbb{E}_{t\sim\text{Unif},x\sim p_t}[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)] +C_1 \end{aligned}$$

令最后和 $\theta$ 无关的一项为常数 $C_1$，然后打开中间这项：

$$\begin{aligned} \mathbb{E}_{t\sim\text{Unif},x\sim p_t}[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)] &= \int_0^1\int p_t(x)u_t^{\theta}(x)^\top u_t^{\text{target}}(x)\mathrm{d} x\mathrm{d} t\\ &= \int_0^1\int p_t(x) u_t^{\theta}(x)^\top \left[\int u_t^{\text{target}} (x|z) \frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}\mathrm{d}z\right] \mathrm{d} x\mathrm{d} t\\ &= \int_0^1\int \int u_t^{\theta}(x)^\top u_t^{\text{target}} (x|z) p_t(x|z)p_{\text{data}}(z)\mathrm{d}z \mathrm{d} x\mathrm{d} t\\ &= \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}}, x\sim p_t(\cdot|z)}[u_t^{\theta}(x)^\top u_t^{\text{target}}(x|z)] \end{aligned}$$

发现我们居然得到了条件向量场相关的项，把这个往回带，然后加一项减一项 $\|u_t^{\text{target}}(x|z)\|^2$ 来配凑：

$$\begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{init}},x\sim p_t(\cdot|z)}\left[\|u_t^{\theta}(x)\|^2\right] - 2\mathbb{E}_{t\sim\text{Unif},x\sim p_t}\left[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)\right] +C_1\\ &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{init}},x\sim p_t(\cdot|z)}\left[\|u_t^{\theta}(x)\|^2- 2u_t^{\theta}(x)^\top u_t^{\text{target}}(x|z) + \|u_t^{\text{target}}(x|z)\|^2 - \|u_t^{\text{target}}(x|z)\|^2\right]+ C_1\\ &= \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x|z) \|^2\right] + \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[-\|u_t^{\text{target}}(x|z)\|^2\right]+C_1\\ &= \mathcal{L}_{\text{CFM}}(\theta)+C_2+C_1 \end{aligned}$$

最后一步是因为 $\mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[-\|u_t^{\text{target}}(x|z)\|^2\right]$ 也是与 $\theta$ 无关的常数。所以推导出来发现 FM loss 和 CFM loss 真的只差一个常数。

这个算法有几个关键特性：

• Simulation-free：训练的时候从来不真正模拟 ODE，所以节省很多资源；
• 训练的目标是一个非常简单的回归目标 $u_t^{\text{target}}(x|z)$；
• 最后直接使用 Euler 方法就可以进行目标生成（从 $p_{\text{init}}$ 采样，然后用 Euler 法模拟 ODE，最后得到 $p_{\text{data}}$ 分布下的样本）。

这就叫做 Flow Matching（流匹配）。

高斯路径

是 denoising diffusion model 使用的概率路径。令 $\alpha_t,\beta_t$ 为 noise schedulers：两个连续可微且单调的关于 $t$ 的函数，且 $\alpha_0=\beta_1=0$，$\alpha_1=\beta_0=1$。则条件概率路径定义如下：

$$p_t(\cdot| z) = \mathcal{N}(\alpha_t z, \beta_t^2I_d)$$

相当于天然满足了 $p_0(\cdot|z) = \mathcal{N}(0,I_d)$，$p_1(\cdot|z) = \delta_z$。并且其使得采样特别简单：

$$z\sim p_{\text{data}},\ \epsilon\sim\mathcal{N}(0,I)\quad\Rightarrow\quad x=\alpha_t z+\beta_t\epsilon\sim p_t$$

这一步很关键：训练时我们只需要数据样本 $z$ 和噪声 $\epsilon$，就能构造任意时刻的“中间状态” $x$。

接下来我们证明高斯路径对应的条件向量场为

$$u_t^{\text{target}}(x|z) = \left(\dot{\alpha}_t - \frac{\dot\beta_t}{\beta_t}\alpha_t \right)z + \frac{\dot\beta_t}{\beta_t}x$$

证明：定义条件流模型 $\psi_t^{\text{target}}(x|z) = \alpha_t z + \beta_t x$，若 $X_t$ 为初始条件 $X_0\sim p_{\text{init}} = \mathcal N(0,I_d)$，然后根据定义

$$X_t = \psi_t^{\text{target}}(X_0|z) = \alpha_t z + \beta_t X_0\sim \mathcal N(\alpha_t z, \beta_t^2 I_d) = p_t(\cdot | z)$$

现在已经说明轨迹是按照条件概率路径分布的了。接下来推导向量场就就行了。根据流的定义：

$$\begin{aligned} \frac{\mathrm{d}}{\mathrm{d} t}\psi_t^{\text{target}}(x|z) &= u_t^{\text{target}}(\psi_t^{\text{target}}(x|z)|z)\\ \dot\alpha_t z + \dot\beta_t z &= u_t^{\text{target}}(\alpha_t z + \beta_t z | z)\\ \dot\alpha_t z + \dot\beta_t \left( \frac{x-\alpha_t z}{\beta_t} \right) &= u_t^{\text{target}}(x|z)&\text{reparameterize }x\to(x-\alpha_tz)/\beta_t\\ \left(\dot{\alpha}_t - \frac{\dot\beta_t}{\beta_t}\alpha_t \right)z + \frac{\dot\beta_t}{\beta_t}x &= u_t^{\text{target}}(x|z) \end{aligned}$$

接下来把上面的结果代入进 CFM loss：

$$\begin{aligned} \mathcal{L}_{\text{CFM}}(\theta) &= \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x|z) \|^2\right]\\ &= \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\left\lVert u_t^{\theta}(x)-\left(\dot{\alpha}_t - \frac{\dot\beta_t}{\beta_t}\alpha_t \right)z - \frac{\dot\beta_t}{\beta_t}x \right\rVert^2\right]\\ &= \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},\epsilon\sim\mathcal N(0,I_d)}\left[\left\lVert u_t^{\theta}(\alpha_t z + \beta_t\epsilon)-\left(\dot\alpha_t z + \dot\beta_t\epsilon \right) \right\rVert^2\right] \end{aligned}$$

如果我们使用 CondOT probability path，即令 $\alpha_t = t, \beta_t = 1-t$ 的话，形式会变得更简单：

$$\mathcal L_{\text{CFM}}(\theta) =\mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},\epsilon\sim\mathcal N(0,I_d)}\left[\| u_t^{\theta}(t z + (1-t)\epsilon) - (z-\epsilon)\|^2 \right]$$

总结

所以，流匹配是什么呢？核心目标是为了学习边缘向量场 $u_t^{\text{target}}$。为了构造之，先定义条件概率路径 $p_t(x|z)$ 使得 $p_0(\cdot|z) = p_{\text{init}}$ 而 $p_1(\cdot|z) = \delta_z$。接下来构造条件向量场 $u_t^{\text{target}}(x|z)$ 使得其对应的流 $\psi_t^{\text{target}}(x|z)$ 满足

$$X_0\sim p_{\text{init}} \implies X_t= \psi_t^{\text{target}}(X_0|z)\sim p_t(\cdot|z)$$

然后定义边缘向量场

$$u_t^{\text{target}}(x) = \int u_t^{\text{target}} (x|z) \frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}\mathrm{d}z$$

这个向量场可以把 $p_{\text{init}}$ 沿着边缘概率路径转变成 $p_{\text{data}}$。

为了学习之，构造一个条件流匹配损失函数

$$\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t\sim \text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x|z) \|^2\right]$$

对于 CondOT 路径的情况，学习的伪代码如下：

$$\begin{array}{ll} 1 & \textbf{Input. } \text{A dataset of samples }z\sim p_{\text{data}}\text{, and neural network }u_t^{\theta} \\ 2 & \textbf{Method. } \\ 3 & \textbf{for } \text{each mini-batch of data}\textbf{ do} \\ 4 & \qquad\text{Sample a data example }z \text{ from the dataset}\\ 5 & \qquad\text{Sample a random time }t\sim\text{Unif} \\ 6 & \qquad\text{Sample noise }\epsilon\sim\mathcal{N}(0,I_d)\\ 7 & \qquad\text{Set }x=tz+(1-t)\epsilon&(\text{General case: }x\sim p_t(\cdot|z))\\ 8 & \qquad\text{Compute loss } \mathcal{L}(\theta) = \|u_t^{\theta}(x) - (z - \epsilon) \|^2&(\text{General case:} =\| u_t^{\theta}(x)-u_t^{\text{target}}(x|z) \|^2)\\ 9 & \qquad\text{Update }\theta\gets \text{grad\_update}(\mathcal{L}(\theta))\\ 10 & \textbf{end for} \end{array}$$

非常之简单优雅。

Score Matching

Score Function

前面的 flow model 的视角是向量场 $u_t(x)$，而 diffusion 采取的是另一个视角：score function。接下来会先给出 score 的定义，然后一步步说明为什么我们在 diffusion 里面需要 score。

假设 $q(x)$ 为某概率分布，那么其 score function 定义为他的对数似然的梯度 $\nabla \log q(x)$，直觉上来看就是对数似然增长最陡峭的方向。回顾条件概率路径和边缘概率路径的相关概念，我们同样可以定义 conditional score function $\nabla\log p_t(x|z)$ 和 marginal score function $\nabla\log p_t(x)$。并且我们也可以得到

$$\nabla\log p_t(x) = \int\nabla \log p_t(x|z)\frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}\mathrm{d}z$$

条件/边缘 score function 的关系和条件/边缘向量场的关系就很像了。上面的式子证明起来也很简单：

$$\nabla\log p_t(x) = \frac{\nabla p_t(x)}{p_t(x)} = \frac{\nabla \int p_t(x|z)p_{\text{data}}(z)\mathrm{d} z}{p_t(x)} = \frac{\int \nabla p_t(x|z)p_{\text{data}}(z)\mathrm{d} z}{p_t(x)} = \int\nabla \log p_t(x|z)\frac{p_t(x|z)p_{\text{data}}(z)}{p_t(x)}\mathrm{d}z$$

对于高斯概率路径 $p_t(x|z)=\mathcal{N}(x;\alpha_t z,\beta_t^2I_d)$，根据高斯的概率密度函数可以算出来条件 score：

$$\nabla\log p_t(x|z) = \nabla \log\mathcal{N}(x;\alpha_tz,\beta_t^2I_d) = -\frac{x-\alpha_tz}{\beta_t^2}$$

注意到，其是关于 $x,z$ 的线性函数，而 $u_t^{\text{target}}(x|z)$ 也是关于 $x,z$ 的线性函数，所以这二者是可以互相转换的：

$$u_t^{\text{target}}(x|z) = a_t\nabla \log p_t(x|z)+b_tx,\quad a_t=\left(\beta_t^2 \frac{\dot\alpha_t}{\alpha_t} - \dot\beta_t\beta_t\right),\quad b_t = \frac{\dot\alpha_t}{\alpha_t}$$

代入积分的式子会发现，对于边缘的情况也成立：

$$u_t^{\text{target}}(x) = a_t\nabla\log p_t(x) + b_t$$

所以对于高斯路径，学习边缘向量场就等价于学习边缘 score function。

SDE 采样

定理 (SDE Extension Trick)：

条件/边缘向量场如前文定义。对于任意扩散系数 $\sigma_t\ge0$，可以构造通过往 ODE 的动力学里面添加随机动力学来构造 SDE：

$$\begin{aligned} &X_0\sim p_{\text{init}} &\mathrm{d} X_t &= {\color{blue}u_t^{\text{target}}(X_t)\mathrm{d} t} + {\color{green}\frac{\sigma_t^2}{2}\nabla \log p_t(X_t) \mathrm{d} t + \sigma_t \mathrm{d} W_t}\\ && &= \left[{\color{blue}u_t^{\text{target}}(X_t)} + {\color{green}\frac{\sigma_t^2}{2}\nabla\log p_t(X_t)} \right]\mathrm{d} t + {\color{green}\sigma_t }\mathrm{d} W_t\\ \implies&X_t\sim p_t \quad (0\le t\le 1) \end{aligned}$$

这告诉我们，如果要添加 $\sigma_t\mathrm{d}W_t$ 这个随机项的话，需要用一个含 score 的项来“校正”这个分布。

这个的证明依赖 Fokker-Planck 方程，作者的数理基础非常薄弱，本文就不展开了。

对于高斯概率路径，这个对应的 SDE 自然就是

$$\mathrm{d} X_t = \left[\left(a_t + \frac{\sigma_t^2}{2} \right)\nabla\log p_t(X_t) + b_t X_t\right]\mathrm{d} t + \sigma_t \mathrm{d} W_t$$

Score Matching

和 flow matching 很像，定义 score 网络 $s_t^{\theta}:\R^d\times[0,1]\to \R^d$。分别定义 score matching loss 和 denoising score matching loss（下面会解释为什么这样叫）：

$$\begin{aligned} \mathcal{L}_{\text{SM}}(\theta) &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\left\lVert s_t^{\theta}(x) - \nabla\log p_t(x) \right\rVert^2\right]\\ \mathcal{L}_{\text{CSM}}(\theta) &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\left\lVert s_t^{\theta}(x) - \nabla\log p_t(x|z) \right\rVert^2\right]\\ \end{aligned}$$

跟 flow matching 一样的叙事，上面的 untractable，下面的有解析解可以算。而他们正好相差常数 $C$（证明方法一模一样），所以只优化下面那个就可以了。

对于高斯路径，代入式子：

$$\begin{aligned} \mathcal{L}_{\text{CSM}}(\theta) &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\left\lVert s_t^{\theta}(x) - \nabla\log p_t(x|z) \right\rVert^2\right]\\ &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},x\sim p_t(\cdot|z)}\left[\left\lVert s_t^{\theta}(x) +\frac{x-\alpha_tz}{\beta_t^2}\right\rVert^2\right]\\ &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},\epsilon\sim\mathcal{N}(0,I_d)}\left[\left\lVert s_t^{\theta}(\alpha_t z + \beta_t\epsilon) +\frac{\epsilon}{\beta_t}\right\rVert^2\right]\\ &= \mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},\epsilon\sim\mathcal{N}(0,I_d)}\left[\frac{1}{\beta_t^2}\left\lVert \beta_t s_t^{\theta}(\alpha_t z + \beta_t\epsilon) +\epsilon\right\rVert^2\right]\\ \end{aligned}$$

那他什么时候取最小值呢，当然是当 $s_t^{\theta}(x) \to -\frac{\epsilon}{\beta_t}$ 的时候了，究其本质，高斯路径下真实的 conditional score 本身就等于“噪声的线性函数”，所以这个网络本质上是在预测噪声，这就解释了为什么这个 loss 叫做 denoising score matching loss。但是注意到当 $\beta_t\to 0$ 的时候（即噪声水平很低的时候）数值会很不稳定。所以 DDPM 直接把 $\beta_t$ 相关这个系数给丢了，然后把 $s_t^{\theta}$ 重参数化为噪声预测网络 $\epsilon_t^{\theta}:\R^d\times[0,1]\to\R^d$：

$$\epsilon_t^{\theta}(x):=-\beta_ts_t^{\theta}(x)\implies\mathcal{L}_{\text{DDPM}}(\theta)=\mathbb{E}_{t\sim\text{Unif},z\sim p_{\text{data}},\epsilon\sim\mathcal{N}(0,I_d)}\left[\left\lVert \epsilon_t^{\theta}(\alpha_t z + \beta_t\epsilon) -\epsilon\right\rVert^2\right]$$

Guidance

我们当然不满足于只在宽泛的 $p_{\text{data}}$ 中采样，希望模型按照我们的 guidance $y$ 来在一个更小的分布 $p_{\text{data}}(z|y)$ 上生成（比如按照提示词）。这一节将讨论相关内容。

假设现在我们的 guidance 是以文字形式，这些文字 $y$ 处于离散的空间 $\mathcal Y$ 中。定义 guided diffusion model 和 guided vector field $u_t^{\theta}(\cdot|y)$ 如下（不使用“条件”是因为容易引起混淆）。这个时候的 $u^{\theta}$ 就应该是 $\R^d\times\mathcal{Y}\times[0,1]$ 了。在 $\sigma_t = 0$ 的情况下就是 flow model，接下来只讨论 flow。

最简单的情况，我们从 $p_{\text{data}}(z,y)$ 这个联合分布里面采样 $z,y$，然后用下面的损失函数训练：

$$\mathcal{L}^{\text{guided}}_{\text{CFM}}(\theta)=\mathbb{E}_{(z,y)\sim p_{\text{data}(z,y)},t\sim\text{Unif},x\sim p_t(\cdot|z)}\left[\lVert u_t^{\theta}(x|y)-u_t^{\text{target}}(x|z) \rVert^2\right]$$

也没啥区别就是说。但是经验表明，这样的生成效果并不够好，可能源于模型欠拟合，没有学到正确的边缘向量场，也有可能是数据质量（图片以及对应的标注不够完美），总之我们需要某种方式强化这个 guidance。现有的 SOTA 生成模型常用的方法称为 classifier-free guidance (CFG)。

Classifier Guidance

在讨论 CFG 之前先讨论一下这个。

为了简便，讨论上文说的高斯路径。我们知道

$$u_t^{\text{target}}(x|y) = a_t\nabla \log p_t(x|y)+b_tx,\quad a_t=\left(\beta_t^2 \frac{\dot\alpha_t}{\alpha_t} - \dot\beta_t\beta_t\right),\quad b_t = \frac{\dot\alpha_t}{\alpha_t}$$

把 $\nabla\log p_t(x|y)$ 用贝叶斯拆一下：

$$\nabla\log p_t(x|y) = \nabla\log\left(\frac{p_t(x)p_t(y|x)}{p_t(y)} \right) = \nabla\log p_t(x) + \nabla\log p_t(y|x)$$

$\nabla_x\log p_t(y) = 0$ 所以扔了。再把它丢回去：

$$u_t^{\text{target}}(x|y) = a_t(\nabla\log p_t(x)+\nabla\log p_t(y|x)) + b_t x = u_t^{\text{target}}(x) + a_t\nabla\log p_t(y|x)$$

这啥意思，说明 guided 向量场等于 unguided 向量场加上一个在引导变量 $y$ 下的对数似然 $\log p_t(y|x)$ 的梯度，直觉上后者在 控制引导的强度，而人们觉得引导强度还不够大所以加一个系数 $w$：

$$\tilde{u}_t(x|y) = u_t^{\text{target}}(x) + w a_t\nabla\log p_t(y|x)$$

这就是 classifier guidance。这个 $\log p_t(y|x)$ 本质上是在加噪数据上的一个分类器，所以可以用监督学习训练。注意到若 $w\ne 1$，则学习到的 $\tilde{u}_t(x|y)$ 不是“真实的”guided 向量场，这是一种 heuristic。

Classifier-free Guidance

Classifier guidance 看似 make sense，但有如下问题：

• 我们需要在训练一个 flow/diffusion 模型的同时训练一个 classifier；
• 如果 $y$ 比较高维，那 classifier 几乎没法训练。

所以 classifier-free guidance 出现了，正如其名，其不需要单独训练一个 classifier。如何做到的呢？先从

$$\nabla\log p_t(x|y) = \nabla \log p_t(x) + \nabla\log p_t(y|x)$$

这个式子说起，把它再次代入 classifier guidance 的式子：

$$\begin{aligned} \tilde{u}_t(x|y) &= u_t^{\text{target}}(x) + w a_t\nabla\log p_t(y|x)\\ &= u_t^{\text{target}}(x) + w a_t(\nabla\log p_t(x|y)-\nabla\log p_t(x))\\ &= u_t^{\text{target}}(x) - (w b_t x + w a_t\nabla\log p_t(x)) + (wb_t x + wa_t\nabla \log p_t(x|y))\\ &= (1-w)u_t^{\text{target}}(x) + w u_t^{\text{target}}(x|y) \end{aligned}$$

相当于把这个 scaled guided vector field 表示成了 $u_t^{\text{target}}(x)$ 和 $u_t^{\text{target}}(x|y)$ 的线性组合。可能你会说这下可以一起训练这两个模型了？但事实上并不需要。很关键的想法是，在标签空间里添加一个 $\varnothing$ 来表示无条件引导的情况，即 $u_t^{\text{target}}(x)$ 可以视为 $u_t^{\text{target}}(x|\varnothing)$。

值得注意的是，这个式子

$$\tilde{u}_t = (1-w)u_t^{\text{target}}(x) + w u_t^{\text{target}}(x|y)$$

对于任意概率路径都成立，不止是高斯路径。利用高斯路径来推导是一种简化。

那么如何训练呢？目前的数据集里面是没有这个所谓的 $\varnothing$ 标签的，但解决方案也很简单，指定一个超参数 $\eta$ 来使得以 $\eta$ 的概率把数据标签替换成 $\varnothing$ 就可以了，最终的目标函数如下：

$$\begin{aligned} \mathcal{L}_{\text{CFM}}^{\text{CFG}}(\theta) &= \mathbb{E}_{\square}\left[\lVert u_t^{\theta}(x|y)-u_t^{\text{target}}(x|z) \rVert^2\right]\\ \square &= (z,y)\sim p_{\text{data}}, t\sim\text{Unif},x\sim p_t(\cdot|z),\text{replace }y\gets\varnothing\text{ with prob. }\eta \end{aligned}$$

注意到，如果我们使用 $w>1$ 的话，我们拿到的 $X_1$ 就不是完全和 $X_1\sim p_{\text{data}}(\cdot|y)$ 对齐了，但经验说明取 $w\ge 4$ 的效果会相当之好。

生成的时候就从 $X_0\sim p_{\text{init}}$ 出发，按照 $\tilde u_t(x|y)$ 这个动力学来走就行了。

Architecture Design & Latent Space

这部分懒得更新了，直接看 lecture notes 吧。