Lec03- 训练生成模型

Recap¶

ConditionalMarginal

Flow Matching¶

我们的目标是 \(u_t^{\theta} \approx u_t^{\mathrm{target}}\)，我们需要定义一个训练目标来优化这个网络的参数 \(\theta\).

Flow Matching Loss¶

给定 \(t \sim \mathcal{U}(0, 1), x \sim p_t(\cdot | z), z \sim p_{data}\)

• \(t\) 是一个均匀分布的随机变量，表示时间
• \(z\) 是一个从数据分布中采样的样本
• \(x\) 是一个从条件概率路径中采样的样本

在上节课中，我们提到这个损失函数是 intractable 的，而遍历所有的样本 \(z\) 也是不现实的，但是有一个很关键的 insight: \(L_{\mathbf{FM}}\) 和 \(L_{\mathbf{CFM}}\) 只相差一个常数 \(C\)，而且这个常数 \(C\) 与训练参数 \(\theta\) 无关，所以最小化后者同样可以达到训练目的。 证明过程在课程笔记 Theorem 18 中。

Conditional Flow Matching Loss:

Training Algorithm¶

t = torch.rand(batch_size) # 1. sample t from uniform distribution
z = data_loader.sample(batch_size) # 2. sample z from data distribution
e = torch.randn(batch_size) # 3. sample noise epsilon from standard normal distribution
pred_u = model(t * z + (1 - t) * e, t) # 4. get predicted vector field from our network u_t(x)
target_u = z - e # 5. get target vector field
loss = torch.mean((pred_u - target_u) ** 2) # 6. compute loss
loss.backward() # 7. backpropagate gradients
optimizer.step() # 8. update parameters for our network u_t(x)

Score Matching¶

边际得分函数 (Marginal Score Function, MSF) 允许将 ODE 扩展为具有任意扩散系数的 SDE. 不同于 Flow Matching，这里我们的网络 \(s_t^{\theta}(x)\) 是近似一个目标得分函数(Score Function)，而不是一个向量场。

Score Matching Loss¶

给定 \(t \sim \mathcal{U}(0, 1), x \sim p_t(\cdot | z), z \sim p_{data}\)

Conditional Score Matching Loss:

在这里我们同样可以证明，\(L_{\mathbf{SM}}\) 和 \(L_{\mathbf{CSM}}\) 只相差一个常数，所以最小化后者同样可以达到训练目的。

Training Algorithm¶

Example: Denoising Score Matching for Gaussian Probability Path

以高斯分布为例：

\[\nabla\log p_t(x|z)=-\frac{x-\alpha_tz}{\beta_t^2}, \epsilon\sim\mathcal{N}(0,I_d) \\ \Rightarrow x=\alpha_tz+\beta_t\epsilon\sim\mathcal{N}(\alpha_tz,\beta_t^2I_d)\]

Denoising Score Matching Loss:

\[\mathcal{L}_\mathrm{dsm}(\theta)= \mathbb{E}_{t\sim\mathrm{Unif},z\sim p_{\mathrm{data}},\epsilon\sim\mathcal{N}(0,I_{d})}[\|s_{t}^{\theta}(\alpha_{t}z+\beta_{t}\epsilon)+\frac{\epsilon}{\beta_{t}}\|^{2}]\]

Training Algorithm:

现在，我们已经设计好了训练算法，训练后可以用 diffusion model 来生成样本了。 我们用一个 SDE 来描述随机采样的过程：

Denoising Diffusion Models¶

去噪扩散模型 (Denoising Diffusion Models, DDMs) 是一种基于扩散模型的生成模型，我们一般直接把它叫做 diffusion model. 去噪扩散模型仅适用于高斯初始分布 \(p_{\mathrm{init}} = \mathcal{N}(0, I_d)\) 和高斯概率路径。
Denoising Diffusion Models = Diffusion Models with Gaussian Probability Path.

DDMs 有很多很好的性质，它的向量场和得分函数都是 \(x\) 和 \(z\) 的加权和，且是线性的，所以可以相互转换。

于是可以将得分函数的神经网络转换为向量场的神经网络后训练：\(u_t^\theta=\left(\beta_t^2\frac{\dot{\alpha}_t}{\alpha_t}-\dot{\beta}_t\beta_t\right)s_t^\theta(x)+\frac{\dot{\alpha}_t}{\alpha_t}x\)
因此我们不需要分别训两个神经网络，最早的 diffusion model 只用 Score Matching 训练了一个得分函数的神经网络。

Diffusion with Conditional Flow Matching¶

在 CondOT Flow Matching 的例子中，我们通过 OT 将条件概率路径设置成连接两个分布采样点的直线：

在 Diffusion 中，更一般的高斯条件概率路径为：

其中 \(\mu_t(x_1)\) 和 \(\sigma_t(x_1)\) 是表示高斯分布的均值和方差的函数。

通过求解以下的 ODE 方程可以得到条件向量场 \(u_t(x|x_1)\):

该向量场的解析解为：

于是，根据不同的微分函数来 \(\mu_t(x_1)\) 和 \(\sigma_t(x_1)\) 我们可以得到不同的条件向量场 \(u_t(x|x_1)\)，注意这里的 condition 也可以是 \((x_0, x_1)\) 的 pair.

• Diffusion conditional Vector Fields:
  • Variance Exploding Schedule (VE): \(\mu_t(x_1) = x_1, \sigma_t(x_1) = \sigma_{1-t}\).
  • Variance Preserving Schedule (VP): \(\mu_t(x_1) = \alpha_{1-t}x_1, \sigma_t(x_1) = \sqrt{1 - \alpha_{1-t}^2}\).
• Optimal Transport conditional Vector Fields (CondOT, most common flow matching schedule)
  • \(\mu_t(x_1) = t x_1, \sigma_t(x_1) = 1 - (1 - \sigma_{min})t\).

将 Diffusion 和 Flow Matching 训练目标结合起来优化，相比于现有方法，训练更加稳定。

评论