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\tag{1}$$

即将 $[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 就是希望我们的轨迹是沿着这个向量场走的：

$$\newcommand{\d}{\mathrm{d}} \begin{aligned}\frac{\d}{\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)\\ \newcommand{\d}{\mathrm{d}} \begin{aligned}\frac{\d}{\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}}\\ \newcommand{\d}{\mathrm{d}} \frac{\d}{\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 转写一下：

$$\newcommand{\d}{\mathrm{d}} \begin{aligned}\frac{\d}{\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)$ 为随机误差项 \ref{1}。上面这个式子就构成一个随机微分方程，可以用下面的符号来表示：

$$\newcommand{\d}{\mathrm{d}} \begin{aligned} \d X_t &= u_t(X_t)\d t + \sigma_t \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，我们也可以构造一个生成式模型，称为扩散模型：

$$\newcommand{\d}{\mathrm{d}} \begin{aligned} X_0&\sim p_{\text{init}} \\ \d X_t &= u_t(X_t)\d t + \sigma_t \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)$​：

$$\newcommand{\d}{\mathrm{d}}X_0\sim p_{\text{init}},\quad \frac{\d X_t}{\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$$

满足

$$\newcommand{\d}{\mathrm{d}}X_0\sim p_{\text{init}},\quad \frac{\d}{\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 表示净流出，所以加负号是净流入）

对于边缘向量场的证明：

$$\newcommand{\d}{\mathrm{d}}\begin{aligned} \partial_t p_t(x)&\overset{(i)}{=}\partial_t \int p_t(x | z)p_{\mathrm{data}}(z)\d z \\ &= \int \partial_t p_t(x | z)p_{\mathrm{data}}(z)\d z \\ &\overset{(ii)}{=} \int -\operatorname{div}\!\left(p_t(\cdot | z)u_t^{\mathrm{target}}(\cdot | z)\right)(x)p_{\mathrm{data}}(z)\d z \\ &\overset{(iii)}{=} -\operatorname{div}\!\left(\int p_t(x | z)u_t^{\mathrm{target}}(x | z)p_{\mathrm{data}}(z)\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)}\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 如下：

$$\newcommand{\E}{\mathbb{E}} \begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \E_{t\sim \text{Unif},x\sim p_t}\left[\| u_t^{\theta}(x)-u_t^{\text{target}}(x) \|^2\right]\\ &\overset{(i)}{=} \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 拆开。

$$\newcommand{\E}{\mathbb{E}} \begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \E_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\theta}(x)-u_t^{\text{target}}(x) \|^2]\\ &= \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]\\ &= \E_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\theta}(x)\|^2] - 2\E_{t\sim\text{Unif},x\sim p_t}[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)] + \E_{t\sim\text{Unif},x\sim p_t}[\|u_t^{\text{target}}(x) \|^2]\\ &= \E_{t\sim\text{Unif},z\sim p_{\text{init}},x\sim p_t(\cdot|z)}[\|u_t^{\theta}(x)\|^2] - 2\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$，然后打开中间这项：

$$\newcommand{\E}{\mathbb{E}}\newcommand{\d}{\mathrm{d}} \begin{aligned} \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)\d x\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] \d x\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 \d x\d t\\ &= \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$ 来配凑：

$$\newcommand{\E}{\mathbb{E}} \begin{aligned} \mathcal{L}_{\text{FM}}(\theta) &= \E_{t\sim\text{Unif},z\sim p_{\text{init}},x\sim p_t(\cdot|z)}\left[\|u_t^{\theta}(x)\|^2\right] - 2\E_{t\sim\text{Unif},x\sim p_t}\left[u_t^{\theta}(x)^\top u_t^{\text{target}}(x)\right] +C_1\\ &= \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)$$

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

$$\newcommand{\d}{\mathrm{d}} \begin{aligned} \frac{\d}{\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