FLOW MATCHING FOR GENERATIVE MODELING

生成建模中的流匹配这是首次提出Flow Matching的论文。“流匹配兼容一类一般的高斯概率路径用于噪声样本和数据样本之间的转换——该路径将现有扩散路径作为特定实例。有趣的是我们发现采用带有扩散路径的FM即Flow Matching能为训练扩散模型提供更稳健和稳定的替代方案。此外流匹配为用其他非扩散概率路径训练CNF连续归一化流打开了大门。一个特别值得关注的例子是使用最优输运OT位移插值来定义条件概率路径。这些路径比扩散路径更高效提供更快的训练和采样速度并实现更好的泛化。”“然而仅限于简单的扩散过程导致抽样概率路径空间相当有限导致训练时间非常长需要采用专门的方法例如Song 等人2020a;Zhang 和 Chen 2022用于高效抽样。”解读 diffusion 的路径通常像是先从数据逐渐扩散成噪声生成时再反向去噪。这个路径未必是最短、最直、最好学的路径。论文甚至指出diffusion path 的采样轨迹可能会 “overshoot 超过”也就是冲过目标后又回调造成不必要的 backtracking而 OT path 保证是更直的轨迹。 直观类比 diffusion 像是你从 A 到 B先绕很多弯路最后慢慢靠近目标 Flow Matching OT 则试图让每个样本更像沿着直线、恒定速度从噪声移动到数据。Flow Matching 做了什么Flow Matching 不从 “加噪扩散过程” 出发而是直接问我想让分布沿某条路径pt从噪声变成数据那么对应的速度场/向量场vt(x) 应该是什么神经网络能不能直接学这个速度场也就是说diffusion 多数时候学的是 score这里的score指的是score ∇_x log p_t(x | x_1) 它表示概率密度上升最快的方向。也就是概率密度的 log 对 x 求梯度。即 ∇ x log⁡ pt(x)然后通过 reverse SDE/ODE 采样Flow Matching 学的是 velocity / vector field也就是某个点在时间 t 应该往哪里流动。进一步的论文提出Conditional Flow Matching解决了直接学整体 marginal vector field边缘矢量场 不可求的问题把路径拆成对每个数据样本的 conditional path然后证明这个条件目标和原始 Flow Matching 目标有相同梯度因此可以无模拟、可扩展地训练 CNF。“然而除了可以通过如去噪分数匹配Vincent 2011高效训练的扩散外尚无可扩展的CNF训练算法。事实上最大似然训练如Grathwohl等2018需要昂贵的数值常微分方程模拟而现有无仿真方法要么涉及不可解积分Rozen等2021要么涉及有偏梯度Ben-Hamu等2022。”解读一下为什么无解diffusion存在无解的情况吗是怎么样的无解fm为什么就能克服这种情况为什么 FM 可以超越 diffusion在 FM 之前diffusion 是少数能大规模训练的连续生成模型因为 denoising score matching 绕开了不可解的整体分布而普通CNF要么训练时解ODE太贵要么无仿真训练目标有不可解积分或有偏梯度。FM 的贡献是把 diffusion 的“条件监督无偏训练”思想推广到 CNF 的 vector field 学习上从而让通用 CNF 也能无仿真、无偏、可扩展地训练。所以不是 diffusion 无解。更准确地说 真实 score / 真实 marginal vector field 不可直接解析求解diffusion 靠 denoising score matching 绕过去FM 靠 Conditional Flow Matching 把这个绕法推广到更一般的 CNF。 因为 diffusion 虽然训练高效但它的路径设计受限通常来自简单扩散过程比如 VP、VE diffusion。FM 不再必须从扩散过程出发而是可以直接设计 probability path。论文中特别强调FM 可以使用 diffusion path也可以使用非 diffusion path比如 Optimal Transport path。 OT path 的好处是更直、更简单。论文中说OT conditional path 的粒子沿直线、恒定速度运动diffusion path 可能会 overshoot导致不必要的回退。OT vector field 的方向在时间上更稳定因此回归任务更简单。所以 FM 克服的不只是“训练 CNF 难”的问题还克服了 diffusion 的一个路径限制diffusion 有高效训练法但路径受限CNF 路径自由但以前训练不高效FM 让 CNF 既路径自由又能高效训练。首先本文先介绍了问题背景p_t(x)随时间变化的概率密度t0 时是简单分布比如高斯噪声 p_0。t1 时希望变成复杂的数据分布 p_1比如图像分布。中间的 p_t 可以理解为“噪声逐渐变成数据”的每一帧分布。v_t(x)随时间变化的向量场它告诉你如果当前某个点在位置 x时间是 t它应该往哪个方向、以多快速度移动。接着是Flow Matching“设x1表示根据某个未知数据分布qx1分布的随机变量。我们假设只能访问 qx1 的数据样本但无法访问密度函数本身。此外设 pt 为一条概率路径使得 p0 p 是一个简单分布例如标准正态分布 px N x|0 I且 p1 分布大致等于 q。我们稍后将讨论如何构建这样的路径。随后设计了流量匹配目标以匹配该目标概率路径使我们能够从p0流向p1。给定目标概率密度路径ptx和对应的矢量场utx生成ptx我们定义流匹配FM目标为”解读FM 是理论上最直接的目标学整体速度场 u_t(x)但整体速度场不可直接算。CFM 是实际可训练的目标用每个数据样本 x_1 的条件速度 u_t(x | x_1) 做监督接着提出CFM“不幸的是由于边际概率路径和VF方程6和8定义中的积分难以解计算ut仍然难以计算因此也难以自然计算原始流量匹配目标的无偏估计量。相反我们提出一个更简单的目标令人惊讶的是这会得到与原始目标相同的最优解。具体来说我们考虑条件流匹配CFM目标”这里相当于基于已知的x1到x的概率值来去预测流。即附录中有所体现基于已知的小尺寸的图片来预测大尺寸的图片或者基于不清晰的图片来预测清晰的图片“4.1 高斯条件概率路径的特殊实例我们的表述对于任意函数μtx1和σtx1完全通用并且可以将它们设置为满足所需边界条件的任何可微函数。我们首先讨论恢复对应于先前扩散过程的概率路径的特殊情况。由于我们直接处理概率路径因此可以完全放弃对扩散过程的推理。因此在下面的第二个例子中我们直接基于Wasserstein-2最优运输解作为一个有趣的实例构建了概率路径。示例 I扩散条件性向量场。扩散模型从数据点开始逐渐添加噪声直到接近纯噪声。这些可以被表述为随机过程对于在任意时间t获得闭式表示有严格要求从而产生高斯条件概率路径ptx|x1并有特定的平均μtx1和std σtx1Sohl-Dickstein等2015;Ho 等2020;Song 等2020b。例如反向噪声→数据方差爆炸VE路径的形式为”“示例二最优运输OT条件VF。条件概率路径中更自然的选择是定义均值和标准值使其随时间线性变化即”解读相当于OT 路径是在用最直接的搬运方式把噪声分布搬到目标数据点附近。所以具体工程实践上需要注意三个层面的事网络架构DiT 或 U-Net训练目标diffusion objective 或 Flow Matching objective概率路径diffusion path 或 OT path它们不是同一类东西。第一层网络架构DiT 和 U-Net 都是“函数近似器”。它们做的事情本质上都是输入x_t, t, 条件信息输出一个和 x_t 同形状的张量比如图像 latent 是x_t shape [C, H, W]模型输出也是output shape [C, H, W]区别只是网络结构不同U-Net卷积结构多尺度下采样/上采样扩散模型里很常见DiTTransformer 结构把图像 latent 切成 patch token用 self-attention 建模所以 DiT 本身不决定你训练什么目标。它只是 backbone。第二层diffusion objective传统 diffusion 常见训练方式是x_1 真实数据ε ~ N(0, I)x_t α_t x_1 σ_t ε模型输入 x_t, t训练它预测噪声ε_θ(x_t, t) ≈ ε损失是L E || ε_θ(x_t, t) - ε ||^2这就是最典型的 diffusion objective。也可以预测 score、x_0、velocity v但核心思想类似 模型学习如何从带噪样本中恢复去噪方向。所以DiT diffusion objective意思是 用 DiT 作为网络让它预测 diffusion 训练目标比如噪声 ε。同理U-Net diffusion objective意思是 用 U-Net 作为网络让它预测 diffusion 训练目标。Stable Diffusion 早期主要是 U-Net diffusion objective。DiT 则是把 U-Net 换成 Transformer。第三层Flow Matching objectiveFlow Matching 不让模型预测噪声而是让模型预测速度场。它构造x_0 噪声x_1 数据x_t 某条路径上的中间点然后训练v_θ(x_t, t) ≈ u_t(x_t | x_1)损失是L_FM E || v_θ(x_t, t) - u_t(x_t | x_1) ||^2这里的 u_t 是“当前点应该怎么移动”的速度。生成时从噪声开始解 ODEdx/dt v_θ(x, t)从 t0 解到 t1得到图像。所以DiT Flow Matching objective意思是 用 DiT 作为网络让它预测速度场。U-Net Flow Matching意思是 用 U-Net 作为网络让它预测速度场。网络可以换训练目标还是 FM。第四层OT path Flow MatchingFlow Matching 还需要指定路径。OT path 是一种特别简单的路径。论文里大概是x_t [1 - (1 - σ_min)t] x_0 t x_1如果忽略很小的 σ_min近似就是x_t ≈ (1 - t)x_0 t x_1也就是从噪声 x_0 直线走到数据 x_1。目标速度是u_t x_1 - (1 - σ_min)x_0近似就是u_t ≈ x_1 - x_0所以DiT OT path Flow Matching意思是 用 DiT 作为网络用 Flow Matching 损失训练路径选 OT 直线路径。完整训练大概是x_1 数据x_0 ~ N(0, I)t ~ Uniform(0, 1)x_t [1 - (1 - σ_min)t] x_0 t x_1target x_1 - (1 - σ_min)x_0loss || DiT(x_t, t) - target ||^2如果换成 U-Net就是loss || U-Net(x_t, t) - target ||^2