Flow Matching and Diffusion Models

1. 概述

生成模型：

• 对象建模为向量 $x∈\mathbb{R}^d$
• 生成即采样 $p_{\mathrm{data}}(x)$
• 数据集 $z_1,z_2,...,z_n~\sim~p_{\mathrm{data}}$
• 条件生成 $p_{\mathrm{data}}(\cdot|y)$

目标：
从易于采样的 $p_{\mathrm{init}}$ (通常 $\sim \mathcal{N}(0,I_d)$ )出发，通过模型转化为 $p_{\mathrm{data}}$ 。

本文是该note的阅读笔记。

2. 模型

Flow Model

通过ODE建模，

$$\begin{aligned} \mathrm{d}X_t &= u_t(X_t)\mathrm{d}t & \text{ODE}\\ X_0 &= x_0 & \text{initial condition} \end{aligned}$$

$X$ 称作轨迹(trajectory)， $u$ 称作向量场(vector field)，而其解 $\varphi_t$ 是 $\mathbb{R}^d \to \mathbb{R}^d$ 的微分同胚，称作流(flow)：

$$\mathrm{d}\varphi_t(X_0) = u_t(\varphi_t(X_0))\mathrm{d}t$$

虽然称作flow model，但是神经网络建模的是向量场 $u^\theta_t$ 而不是 $\varphi$ 。

Diffusion Model

通过SDE建模，

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

通过Euler-Maruyama method离散求解：

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

当 $\sigma=0$ 时，变为flow model。

SDE生成模型：

• 神经网络： $u^\theta:\mathbb{R}^d\times [0,1]\to\mathbb{R}^d,(x,t)\mapsto u_t^\theta(x)$
• 固定的扩散系数： $\sigma_t:[0,1]\to[0,\infty),t\mapsto\sigma_t$

采样：

• 初始化： $X_0\sim p_{\mathrm{init}}$
• 模拟： $\mathrm{d}X_t=u_t^\theta(X_t)\mathrm{d}t+\sigma_t\mathrm{d}W_t$
• 目标： $X_1\sim p_{\mathrm{data}}$

3. 构建训练目标

为了训练一个神经网络，我们还需要一个损失函数，比如均方误差(mean-squared error):

$$\mathcal{L(\theta)=||u^\theta_t(x)-\underbrace{u_t^{\mathrm{target}}(x)}_{\text{training target}}||^2}$$

因此接下来需要找到描述训练目标的（估计）公式。

3.1. 条件和边缘概率路径（Conditional and Marginal Probability Path）

一个条件（插值）概率路径长这样：

$$\begin{aligned} & p_0(\cdot|z)= p_{\mathrm{init}}, & p_1(\cdot|z) = \delta_{z} \end{aligned}$$

换言之，条件概率路径会逐步将单个数据点转化为初始分布。

边缘概率路径：

$$z\sim p_{\mathrm{data}},~~~x\sim p_t(\cdot|z)~~~\Rightarrow x\sim p_t ~~~（从p_t中采样）\\ p_t(x)=\int p_t(x|z)p_{\mathrm{data}}(z)\,\mathop{}\!\mathrm{d}z ~~~（难以计算）$$

显然：

$$p_0=p_{\mathrm{init}},~p_1=p_{\mathrm{data}}$$

Example: 高斯条件概率路径（Gaussian Conditional Probability Path）

$p_t(\cdot|z)\sim\mathcal{N}(\alpha_tz,\beta_t^2I_d)~~~~~~\alpha_0=\beta_1=0,\alpha_1=\beta_0=1$

从高斯条件路径中采样：

$z\sim p_{\mathrm{data}},\,\epsilon\sim\mathcal{N}(0,I_d)\Rightarrow x=\alpha_tz+\beta_t\epsilon\,\sim p_t$

3.2. 条件和边缘向量场（Conditional and Marginal Vector Fields）

未完待续