Palind's Blog
0%

Flow Matching（流匹配）学习笔记

发表于 更新于 分类于 ，
本文字数： 9.8k 阅读时长 ≈ 9 分钟

在开始学习 Flow Matching 之前，我需要回顾以前的学习内容，并明确学习目的。

回顾

有这样的几类主要生成模型：

  • VAE（变分自编码器，Variational Auto-encoder）。VAE 通过编码器将输入数据映射到潜在变量$z$的潜在空间（latent space），再用解码器重建原始输入。
    输入$x$对应的真实后验分布$p(z∣x)$是潜在变量$z$的分布，是一个连续的多维概率分布，人们常常假设它为高斯分布（因此，人们设置训练的近似分布$q_ϕ(z∣x)$是高斯分布）。编码器想学习出来的，就是这个分布，或者说，就是输入对应的潜在变量$z$的分布参数（高斯分布的均值、方差），而不是$z$自己。我的学习笔记是VAE学习笔记 | Palind’s Blog

    为什么不是$z$自己？因为如果学习的是$z$自己，模型就退化为了一个常见的 Auto-encoder，潜在空间是不连续的。
    VAE 名字中的“Variational”来自变分推断（Variational Inference），通过训练变分分布$q_ϕ(z∣x)$来接近真实的$p(z∣x)$后验。
    学习后验分布，使得模型的潜在空间是连续的。模型泛化能力更强，编码器输出的是分布，那我们就可以从中采样$z$输入解码器，使模型生成多样性的输出。此外，还可以进行插值变换。

  • GAN（生成对抗网络，Generative Adversarial Network）。它由生成器（Generator）和判别器（Discriminator）组成。Generator 从随机噪声$z$中生成样本，Discriminator 判断输入是真实样本还是生成样本。一起训练后，Generator 就可能生成很逼真的样本。我的学习笔记是GAN学习笔记 | Palind’s Blog

  • Diffusion Model（扩散模型）。它通过逐步添噪的扩散过程，学习逆过程逐步去噪，以此从纯噪声中逐渐恢复出清晰的生成样本。我的学习笔记（学的是DDPM）是Diffusion Model学习笔记 | Palind’s Blog

  • Flow-based Method（流式生成模型）。Flow-based 与 VAE 很像。既然 VAE 的 Encoder 和 Decoder 的行为是相反的，一个是从图片生成潜在变量的分布，一个是从分布的采样（也就是潜在变量）生成回图片，那如果训练一个可逆的（invertible） 模型，就可以做生成和反推的工作。我的学习笔记是Flow based Generative Model学习笔记 | Palind’s Blog

Lil’Log 上有一张很经典的图，我现在把它放在笔记里便于我回顾：

1

参考

我会在笔记中放上这些参考的截图，为方便我学习与记忆。

Continuous Normalizing Flow（CNF）及其遇到的问题

回顾 Residual flows，

如果层数较多，那么每一次的变换函数 u 进行的改变就很小。我们可以把这个“变换很小”用$\delta$人为显式表示出来：

$\delta$ 表示 u 的作用强度，移项，除以$\delta$ ，

当层数非常多时，每一层的更新都是对某个连续动力系统的微小近似，

这是 Ordinary Differential Equation（ODE，常微分方程）。它的含义是：样本位置的变化，遵循一个向量场。

那我们的目标就可以是，用一个带有参数 θ 的神经网络，来表示这个参数化的向量场，这样我们就知道样本是怎么从高斯分布一步一步变为我们想要的目标分布的了：

这就是 Neural Ordinary Differential Equation（Neural ODE），简单解释一下，下标 t 代表时刻，这个向量场随时间变化。

以下是一个一维的高斯分布 在时间的变化下随着一维的向量场变化为一维目标分布 的例子：

2

对于某一个时刻 t，特定位置的概率密度是如何变化的呢？我们接下来推导。不过，在推导之前，最好解释一下为什么我们要推导特定位置的概率密度：因为我们训练 CNF 的目标是最大化训练样本在模型分布下的对数似然，所以如果不求它，就没法最大化 log-likelihood，也找不到 loss 函数。

现在来推导。我们以一维为例。在时间 t 时，对于位置$\mathbf{x}’$，对应概率密度$p_t(\mathbf{x}’)$。当 t 加上 $\epsilon$ 后，概率密度变为$p_{t+\epsilon}(\mathbf{x}’)$。

如何用向量场表达这一点的概率密度从$p_t(\mathbf{x}’)$到$p_{t+\epsilon}(\mathbf{x}’)$的变化？

向量场一定是把样本退离或推向了这个位置。为什么会变？因为样本从别的地方流进来了，或者从这里流走了。

2

$u_t(\mathbf{x},θ)$ Vector field（上图左下角的式子）：这是一个向量场。它表示在时刻 t 时，在空间中的任何位置 $\mathbf{x}$ 处，样本将要移动的瞬时速度和方向。

$\frac{d}{dx}$ (或说是 ∇⋅)：这个符号表示对向量场 $u_t(\mathbf{x},θ)$ 求散度（Divergence）。

散度是一个标量场，它度量了在某个点上向量场向外扩散（outflow）或向内收缩（inflow）的程度。我们还是以一维的为例：

如果 $\frac{d}{dx}>0$：

  • 这表示当沿着 x 的正方向移动时，速度 u 在增加。举例：在 x 处速度是 2 米/秒，在 x+Δx 处速度是 2.5 米/秒。
    这意味在 x 附近，右侧的流速比左侧的流速要快。样本会从左侧（流速较慢）向右侧（流速较快）扩散，导致在点 x 附近的样本密度会减小。这个区域就会被“拉开”，里面的东西就会变得稀疏，就是图片中说的“Local outgoingness”（局部出流）。

如果 $\frac{d}{dx}<0$：

  • 这表示当沿着 x 的正方向移动时，速度 u 在减小。这意味着，在 x 附近，右侧的流速比左侧的流速要慢。因此，样本会从右侧（流速较快）追上左侧（流速较慢），导致在点 x 附近的样本密度会增加。

总结一下，$u_t(\mathbf{x},θ)$ (Vector field)描述了在时刻 t，每个位置 x 处样本的瞬时移动速度和方向。它驱动着样本从一个分布流向另一个分布。$\frac{d}{dx}$ (散度)描述了在时刻 t，位置 x 处样本密度变化的趋势。它不是直接的速度，而是速度在空间上的变化率。

$p_t(x_t)u_t(x_t)$：概率流（Probability Current）或概率通量（Probability Flux）。它表示在时刻 t，通过位置 $x_t$ 的“样本流量”或者“概率流量”的大小和方向。即单位时间内，有多少“概率”带着速度 $u_t(x_t)$ 流过 $x_t$ 这个点。

$\frac{d}{dx}(p_t(\mathbf{x}_t)u_t(\mathbf{x}_t))$：这是“概率流”的散度，表示概率流在空间上的变化率。

必然有：

一维的情况就是这样，在高维的空间中，是：

移项，

这就是Continuity Equation（连续性方程）或者叫 transport equation。

为了得到对数似然，以最大化似然函数，我们现在继续推导。

现在考虑对数概率密度 $L = \log p_t(\mathbf{x}_t)$。

对其求全导数（物质导数/随体导数）：

这里补充一下全导数（物质导数），它是流体微团所携带的物理量随时间的变化率，
它由直接对时间的偏导数（固定点 x 不动，这个空间点上的物理量 ϕ 随时间）与对流项相加计算。
对流项就是对空间用链式法则：

代入物质导数方程：

回到展开后的连续性方程，重新排列：

现在我们有了

就可以计算对数似然，本质就是联合概率：

逆变换（Inverse Transform）就是：

但是，在我们计算对数似然的时候，后一项要求解常微分方程（ODE）的积分，限制了性能。

我们用Flow Matching（流匹配）来实现Continuous Normalizing Flow，简化训练过程，消除对昂贵的 ODE 求解的依赖。

Flow Matching

Flow Matching Objective

回顾，连续方程：

刚刚提到过，计算概率密度过于昂贵，我们不让机器去学习如何计算概率密度，而是让它学会匹配（match）流（flow），就是匹配$u_t(\mathbf{x}_t)$。这就是流匹配。

随时间变化的向量场$u_t$完全决定了概率路径和最终目标的分布（target distribution）。既然目的是训练网络来匹配$u_t(\mathbf{x}_t)$（神经网络去拟合某个向量场的回归问题），那么 loss 函数就是：

当 t = 0 时，为$p_{\text{base}}(\mathbf{z})$正态分布，当 t = 1时，为目标分布 $p_{\theta}(\mathbf{x})$，如图：

4

既然 loss 函数是二范数损失（L2 Loss），那就很容易训练。

但是，我们不知道上图中的概率路径（Probability path）和向量场，就是 loss 函数公式中绿色与蓝色的部分。

构造概率路径（Probability path）

我们需要用条件概率的传递构造概率路径，定义时间 t 的概率分布，创造训练数据：

这里，$q(\mathbf{z})$就是$p_{\text{base}}(\mathbf{z})$。我们的边缘概率路径现在被条件概率路径表示，它随着条件变量$\mathbf{z}$的变化而变化。

这样，我们接下来只需要设计一个有效的条件概率路径和向量场就行了。

设计有效的条件概率路径（Conditioinal probability path）和向量场，得到 Conditional Flow Matching (CFM) Objective

Example 1

假设我们的条件是训练数据集中的单个数据点，

条件概率路径就是：

t = 0 时，为标准正态分布；t = 1 时，分布移动到了 $\mathbf{x}_1$。

向量场也很简明：

表达的就是“朝向$\mathbf{x}_1$处走”。速度随时间变化，保证t = 1时，准确走到$\mathbf{x}_1$处。

为什么这个向量场的设计能保证在 $t = 1$ 时准确走到 $\mathbf{x}_1$？

我们给了一个 时间依赖的向量场：

初始时刻 $t = 0$，我们从某个 $\mathbf{x}(0) = \mathbf{x}_0$ 出发。这个方程的解 $\mathbf{x}(t)$ 会不会在 $t = 1$ 时正好变成 $\mathbf{x}_1$ ？我们要解这个常微分方程（ODE）。

为了方便书写，我们先设：

  • 令 $a = 1 - \sigma$，那么 $0 < a < 1$

所以原方程变成：

我们把它写成常见形式：

这是同济的《高等数学》教过的 一阶线性微分方程，直接算出来：

$\mathbf{x}(0) = \mathbf{x}_0$，代入：

所以最后解是：

验证 $t = 1$ 时是否走到 $\mathbf{x}_1$，代入 $t = 1$：

我们的 loss 函数就是：

无论有无条件，可以证明：

这提供了一些可扩展的训练方式。

Example 2：Independent Coupling（独立耦合）

还有一些其他的设计。比如，不仅以目标数据点 $\mathbf{x_1}$为条件，还从 base distribution 中采样一个noise：$\mathbf{x}_0$，作为路径的起点。

条件概率路径就可以是（小方差的）从$\mathbf{x}_0$移动到$\mathbf{x}_1$的高斯分布。

条件向量场是常量。如图：

5

Example 3：Optimal transport CFM

也可以从 base distribution 中抽取多个样本，与训练数据集中的多个数据点建立对应关系。 如图：

6

对比一下三种设计：

7

训练过程

如下图，首先，根据设计，采样数据点$\mathbf{x}_1$与噪声$\mathbf{x}_0$。根据概率路径，创建样本$\mathbf{x}_t$。用此样本训练网络学习向量场，其 loss 函数是$\mathcal{L}_\text{CFM}$。

8

与 Diffusion Model 对比：

9

全文完，感谢阅读😊

作者：Palind