【生成式模型入门】(四):Flow Matching——从速度场学习到生成路径

前言


文章目录

    • 前言
    • [1 扩散模型的局限](#1 扩散模型的局限)
        • [1-1 回顾 Diffusion Policy](#1-1 回顾 Diffusion Policy)
        • [1-2 迭代去噪的计算代价](#1-2 迭代去噪的计算代价)
    • [2 Flow Matching](#2 Flow Matching)
        • [2-1 介绍](#2-1 介绍)
          • [2-1-1 思路提出](#2-1-1 思路提出)
          • [2-1-2 总览](#2-1-2 总览)
        • [2-2 整体框架](#2-2 整体框架)
          • [2-2-1 总览](#2-2-1 总览)
          • [2-2-3 目标](#2-2-3 目标)
        • [2-3 速度场与常微分方程(ODE)](#2-3 速度场与常微分方程(ODE))
          • [2-3-1 ODE](#2-3-1 ODE)
          • [2-3-2 连续性方程](#2-3-2 连续性方程)
          • [2-3-3 代码](#2-3-3 代码)
        • [2-4 条件流匹配(Conditional Flow Matching)](#2-4 条件流匹配(Conditional Flow Matching))
          • [2-4-1 边缘速度场与条件速度场](#2-4-1 边缘速度场与条件速度场)
          • [2-4-2 条件流匹配的核心思路](#2-4-2 条件流匹配的核心思路)
          • [2-4-3 构造高斯条件路径](#2-4-3 构造高斯条件路径)
          • [2-4-4 构造线性插值路径](#2-4-4 构造线性插值路径)
            • [2-4-4-1 CondOT(Conditional Optimal Transport,条件最优传输)](#2-4-4-1 CondOT(Conditional Optimal Transport,条件最优传输))
            • [2-4-4-2 最优传输](#2-4-4-2 最优传输)
            • [2-4-4-3 代码实现](#2-4-4-3 代码实现)
            • [2-4-4-4 从路径速度到条件速度场](#2-4-4-4 从路径速度到条件速度场)
          • [2-4-5 从条件速度到边缘速度](#2-4-5 从条件速度到边缘速度)
          • [2-4-6 Flow Matching Loss](#2-4-6 Flow Matching Loss)
        • [2-5 训练与推理流程](#2-5 训练与推理流程)
          • [2-5-1 训练:随机采样时间并拟合速度](#2-5-1 训练:随机采样时间并拟合速度)
          • [2-5-2 推理:从噪声积分到目标样本](#2-5-2 推理:从噪声积分到目标样本)
        • [2-6 Rectified Flow:让生成路径更直接](#2-6 Rectified Flow:让生成路径更直接)
    • [3 Flow Matching 网络架构与实现](#3 Flow Matching 网络架构与实现)
        • [3-1 整体架构](#3-1 整体架构)
        • [3-2 时间嵌入与条件输入](#3-2 时间嵌入与条件输入)
        • [3-3 速度预测网络](#3-3 速度预测网络)
        • [3-4 训练代码解析](#3-4 训练代码解析)
        • [3-5 ODE 数值积分与采样代码](#3-5 ODE 数值积分与采样代码)
        • [3-6 采样步数与生成质量](#3-6 采样步数与生成质量)
    • [4 实验:从噪声生成数据](#4 实验:从噪声生成数据)
        • [4-1 二维合成数据集](#4-1 二维合成数据集)
        • [4-2 训练过程与损失曲线](#4-2 训练过程与损失曲线)
          • [4-2-1 损失平了不等于质量平了](#4-2-1 损失平了不等于质量平了)
        • [4-3 训练并生成双月形数据](#4-3 训练并生成双月形数据)
        • [4-4 速度场与生成轨迹](#4-4 速度场与生成轨迹)
          • [4-4-1 t 接近 1 时速度场为什么发散](#4-4-1 t 接近 1 时速度场为什么发散)
        • [4-5 改变 ODE 积分步数 ------ 少步采样换来多少质量、省下多少时间](#4-5 改变 ODE 积分步数 —— 少步采样换来多少质量、省下多少时间)
          • [4-5-1 单步采样:代价最低,误差最大](#4-5-1 单步采样:代价最低,误差最大)
          • [4-5-2 质量在哪里饱和](#4-5-2 质量在哪里饱和)
          • [4-5-3 代价是线性的,收益不是](#4-5-3 代价是线性的,收益不是)
    • 总结

1 扩散模型的局限

1-1 回顾 Diffusion Policy
  • 第三期我们讲过 Diffusion Policy 的做法:把机器人策略本身写成一个条件扩散模型,以当前观测为条件,从随机噪声出发逐步去噪出一段动作序列
  • 反向过程与第二期的 DDPM 完全一致,只是在每一步多了一路条件输入

x t − 1 = 1 α t ( x t − 1 − α t 1 − α ˉ t   ε θ ( x t , t , c ) ) + σ t z x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}}\, \varepsilon_\theta(x_t, t, c) \right) + \sigma_t z xt−1=αt 1(xt−1−αˉt 1−αtεθ(xt,t,c))+σtz

  • 以防你忘记:
    • x t x_t xt:第 t t t 步的加噪动作序列,也就是被去噪的对象
    • c c c:条件,在 Diffusion Policy 里是当前观测与本体状态
    • ε θ ( x t , t , c ) \varepsilon_\theta(x_t, t, c) εθ(xt,t,c):网络给出的噪声估计
    • σ t z \sigma_t z σtz:每步注入的随机噪声, z ∼ N ( 0 , I ) z \sim \mathcal{N}(0, I) z∼N(0,I)
    • t t t:反向链的步数下标,从 T T T 数到 1
  • 这套写法解决的是第三期最核心的那个问题:多模态的动作分布
    • 绕杯子时从左侧过和从右侧过都是合法解,单步回归策略输出的是二者的折中,也就是一条撞上杯子的轨迹
    • 扩散策略不输出折中解,而是从整个动作分布里采样,因此两条路都能走出来

说人话:扩散策略的价值在于"它会撒骰子",代价也在于"它得撒很多次骰子"

  • 而"撒很多次"就是本章要算的账:上面那个式子只是一步 。要拿到一段能执行的动作,得把 t = T , T − 1 , ... , 1 t = T, T-1, \ldots, 1 t=T,T−1,...,1 全部走一遍
1-2 迭代去噪的计算代价
  • 上面那个式子里,下标 t t t 从 T T T 递减到 1 1 1,而 T T T 通常取值为 1000。这意味着,生成一批动作可能需要调用 1000 次网络前向
    • 单次前向的计算开销并不高,但重复上千次,累计耗时就相当可观
  • 减少反向去噪步数可以加速推理,但步数并非越少越好
    • 步数太多,推理开销大;步数太少,离散求解误差又可能损害生成质量
    • 因此,如何平衡生成质量与推理速度,是一个关键问题
  • 这一问题在机器人任务上尤为突出,因为策略推理受到控制周期的约束
    • 假设控制频率为 10~50 Hz,一个控制周期只有 20~100 ms
    • 对于采用动作块(Action Chunking)的策略,每批动作执行完毕后都需要重新规划。推理必须在有限的时间窗口内完成,上千次网络前向显然难以满足这一要求
  • 那么,为什么生成过程需要反复调用网络?
    • 在 DDPM 中,每一步都需要根据当前的带噪样本估计噪声,再结合预设的噪声调度计算下一步样本。网络负责预测噪声,噪声调度决定更新规则。 整个生成过程被拆分为一系列离散的去噪步骤,每一步都需要重新进行网络推理
  • 这就引出了一个更根本的问题:能否直接学习从噪声走向目标数据的连续变化规律,而不是拘泥于逐步去噪的生成方式? 这正是 Flow Matching 的切入点

2 Flow Matching

2-1 介绍
2-1-1 思路提出
  • Flow Matching 由 Lipman 等人在论文《Flow Matching for Generative Modeling》中提出,核心目标是解决 CNF(Continuous Normalizing Flow,连续归一化流)的训练难题
  • CNF 可以理解为:用一个连续变化的速度场,把简单的噪声分布变换成复杂的数据分布。它的基本形式是

d x ( t ) d t = v θ ( x ( t ) , t ) \frac{dx(t)}{dt}=v_\theta(x(t),t) dtdx(t)=vθ(x(t),t)

  • 传统 CNF 通常采用最大似然训练,需要计算数据的对数似然。除了求解 ODE,还需要计算速度场的散度(divergence),即速度场在局部的膨胀或收缩程度

∇ ⋅ v θ = ∑ i ∂ v θ , i ∂ x i \nabla\cdot v_\theta =\sum_i\frac{\partial v_{\theta,i}}{\partial x_i} ∇⋅vθ=i∑∂xi∂vθ,i

  • 这个计算涉及对速度场求导,通常还要与 ODE 求解过程结合,因此训练代价可能很高
  • Flow Matching 的思路由此展开:既然最终目的是学习一个能够驱动样本从噪声分布走向数据分布的速度场,能否绕开训练时反复求解 ODE 的过程,直接学习这个速度场?
    • 它通过构造概率路径及其对应的目标速度,将训练转化为速度场回归问题
    • 这样,训练时只需计算预测速度与目标速度之间的误差,并通过反向传播更新网络参数,无需为了构造训练目标而反复求解完整的 ODE
    • 到了生成阶段,再从初始噪声出发,通过数值积分求解学习到的速度场,得到最终样本
  • 论文进一步证明,在相应条件下,这种可计算的条件流匹配(CFM,Conditional Flow Matching)目标,与原始 Flow Matching 目标具有相同的参数梯度(见 2-4-6)。因此,我们无需在训练时求解完整的生成 ODE,就能优化所需的速度场

说人话:Flow Matching 的关键,是把生成模型的训练转化为速度预测问题;训练时学习往哪里走,生成时再沿着学到的速度场积分前进。

2-1-2 总览
  • 本章按论文的推导顺序组织:先定义流与速度场(2-3),再构造条件路径并推出可算的目标速度(2-4),然后是训练与推理的流程(2-5)
  • 本章最后处理一个绕不开的问题:它最出名的一个后继工作 Rectified Flow 改了什么,又和 CondOT 差在哪(2-6)
  • 配套的官方实现是 Meta 开源的 flow_matching 库,版本 1.0.10,它同时发布了一份长文指南《Flow Matching Guide and Code》。本文的推导以论文原文为准,引用时按推导内容指认,不另编公式号;第三章逐段对照的代码片段、以及 3-5 关于求解器求值次数的结论,都以该版本的实现为准
2-2 整体框架
2-2-1 总览

#mermaid-svg-fenElimfzB61obsI{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-fenElimfzB61obsI .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-fenElimfzB61obsI .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-fenElimfzB61obsI .error-icon{fill:#552222;}#mermaid-svg-fenElimfzB61obsI .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-fenElimfzB61obsI .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-fenElimfzB61obsI .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-fenElimfzB61obsI .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-fenElimfzB61obsI .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-fenElimfzB61obsI .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-fenElimfzB61obsI .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-fenElimfzB61obsI .marker{fill:#333333;stroke:#333333;}#mermaid-svg-fenElimfzB61obsI .marker.cross{stroke:#333333;}#mermaid-svg-fenElimfzB61obsI svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-fenElimfzB61obsI p{margin:0;}#mermaid-svg-fenElimfzB61obsI .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-fenElimfzB61obsI .cluster-label text{fill:#333;}#mermaid-svg-fenElimfzB61obsI .cluster-label span{color:#333;}#mermaid-svg-fenElimfzB61obsI .cluster-label span p{background-color:transparent;}#mermaid-svg-fenElimfzB61obsI .label text,#mermaid-svg-fenElimfzB61obsI span{fill:#333;color:#333;}#mermaid-svg-fenElimfzB61obsI .node rect,#mermaid-svg-fenElimfzB61obsI .node circle,#mermaid-svg-fenElimfzB61obsI .node ellipse,#mermaid-svg-fenElimfzB61obsI .node polygon,#mermaid-svg-fenElimfzB61obsI .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-fenElimfzB61obsI .rough-node .label text,#mermaid-svg-fenElimfzB61obsI .node .label text,#mermaid-svg-fenElimfzB61obsI .image-shape .label,#mermaid-svg-fenElimfzB61obsI .icon-shape .label{text-anchor:middle;}#mermaid-svg-fenElimfzB61obsI .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-fenElimfzB61obsI .rough-node .label,#mermaid-svg-fenElimfzB61obsI .node .label,#mermaid-svg-fenElimfzB61obsI .image-shape .label,#mermaid-svg-fenElimfzB61obsI .icon-shape .label{text-align:center;}#mermaid-svg-fenElimfzB61obsI .node.clickable{cursor:pointer;}#mermaid-svg-fenElimfzB61obsI .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-fenElimfzB61obsI .arrowheadPath{fill:#333333;}#mermaid-svg-fenElimfzB61obsI .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-fenElimfzB61obsI .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-fenElimfzB61obsI .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-fenElimfzB61obsI .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-fenElimfzB61obsI .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-fenElimfzB61obsI .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-fenElimfzB61obsI .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-fenElimfzB61obsI .cluster text{fill:#333;}#mermaid-svg-fenElimfzB61obsI .cluster span{color:#333;}#mermaid-svg-fenElimfzB61obsI div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-fenElimfzB61obsI .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-fenElimfzB61obsI rect.text{fill:none;stroke-width:0;}#mermaid-svg-fenElimfzB61obsI .icon-shape,#mermaid-svg-fenElimfzB61obsI .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-fenElimfzB61obsI .icon-shape p,#mermaid-svg-fenElimfzB61obsI .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-fenElimfzB61obsI .icon-shape .label rect,#mermaid-svg-fenElimfzB61obsI .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-fenElimfzB61obsI .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-fenElimfzB61obsI .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-fenElimfzB61obsI :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 推理阶段 Inference
采样高斯噪声 x₀
ODE 数值积分

dx/dt = vθ(xₜ,t)
训练好的速度网络
生成样本 x₁
训练阶段 Training
高斯噪声 x₀
线性插值

xₜ = (1-t)x₀ + tx₁
真实样本 x₁
随机时间 t
速度预测网络

vθ(xₜ, t)
目标速度

ẋₜ = x₁ - x₀
Flow Matching Loss

L = ‖vθ - uₜ‖²
反向传播与参数更新

  • 上图把整个流程分成两半。左半是训练,右半是推理,两边的输入并不对称

    • 训练阶段 :从标准高斯分布采样源样本 x 0 x_0 x0,从数据集采样目标样本 x 1 x_1 x1,并从 0 , 1 0,1 0,1 均匀采样时刻 t t t。
    • 推理阶段 :不再需要真实样本 x 1 x_1 x1 和目标速度,只需从噪声 x 0 x_0 x0 出发,利用网络预测速度,通过 ODE 数值积分逐步更新状态,最终生成样本。
  • 整个框架的核心是训练速度网络。线性插值路径预先确定,目标速度可直接计算,推理时再通过数值积分生成样本。
2-2-3 目标
  • 核心我们要网络学习的内容就一个:给它当前位置 x x x 和时刻 t t t,它给出这一刻该走的速度 v θ ( x , t ) v_\theta(x,t) vθ(x,t)

  • 训练时,每一对随机配对都提供一个监督信号 x t = ( 1 − t ) x 0 + t x 1 ⏟ 当前位置 , x ˙ t = x 1 − x 0 ⏟ 目标速度 \underbrace{x_t=(1-t)x_0+tx_1}{\text{当前位置}},\qquad \underbrace{\dot x_t=x_1-x_0}{\text{目标速度}} 当前位置 xt=(1−t)x0+tx1,目标速度 x˙t=x1−x0

  • 网络见的是 x t x_t xt 与 t t t,要预测的是这一对起点终点给出的速度。这样的样本多了,它就学到了"在这个位置、这个时刻该往哪走"

  • 推理时不再需要真实样本,从噪声出发,反复按学到的速度场更新位置

d x d t = v θ ( x , t ) \frac{dx}{dt}=v_\theta(x,t) dtdx=vθ(x,t)

  • 一路积分到 t = 1 t=1 t=1,就得到生成结果

说人话:训练时用大量随机配对教网络每个位置、每个时刻该往哪走;生成时从噪声出发,沿着网络预测的速度一路走到终点,就得到生成结果。

2-3 速度场与常微分方程(ODE)
  • 直觉上,我们可以认为速度场能够推动样本分布发生变化。但要理解这种变化为何成立,以及速度场与概率分布之间存在怎样的数学关系,就需要从 ODE 和连续性方程出发进行分析。
2-3-1 ODE
  • Flow Matching 的数学起点是一个常微分方程(ODE)。定义一族随时间变化的映射 ϕ t \phi_t ϕt( t ∈ 0 , 1 t \in 0,1 t∈0,1),它由下式给出 d d t ϕ t ( x ) = u t ( ϕ t ( x ) ) , ϕ 0 ( x ) = x \frac{d}{dt}\phi_t(x) = u_t(\phi_t(x)), \qquad \phi_0(x) = x dtdϕt(x)=ut(ϕt(x)),ϕ0(x)=x其中, u t u_t ut 是速度场,描述每个位置在时刻 t t t 的瞬时速度; ϕ t ( x ) \phi_t(x) ϕt(x) 表示样本从初始位置 x x x 出发,运动到时刻 t t t 后的位置。
  • 将初始分布 p 0 p_0 p0 通过映射 ϕ t \phi_t ϕt 推前,就得到时刻 t t t 的分布:: p t = ( ϕ t ) ∗ p 0 p_t=(\phi_t)_*p_0 pt=(ϕt)∗p0直观来说, ϕ t \phi_t ϕt 描述样本从初始状态移动到时刻 t t t 的位置,而 p t p_t pt 描述这些样本在此时如何分布。
  • 根据变量替换公式,其概率密度为: p t ( x ) = p 0 ( ϕ t − 1 ( x ) ) ∣ det ⁡ ∂ ϕ t − 1 ( x ) ∂ x ∣ p_t(x)=p_0(\phi_t^{-1}(x)) \left|\det\frac{\partial\phi_t^{-1}(x)}{\partial x}\right| pt(x)=p0(ϕt−1(x)) det∂x∂ϕt−1(x)
  • 直接计算这一表达式需要求逆映射及其雅可比行列式,计算成本较高。
  • 因此,Flow Matching 转而研究速度场与概率密度之间的关系,通过连续性方程描述分布如何随速度场演化,从而学习速度场来实现分布变换,而不必在训练时显式计算上述推前密度公式。
2-3-2 连续性方程
  • 上一节提到,直接通过流映射计算概率密度,需要求逆映射及其雅可比行列式。那么,能否绕开这一步,直接描述概率分布随时间的变化?
    • 直观来说,当速度场推动样本向四周散开时,局部概率密度通常会降低;当样本向某个区域聚集时,局部概率密度则可能升高。
  • 因此,根据概率质量守恒,可以建立概率密度 p t p_t pt 与速度场 u t u_t ut 之间的关系,即连续性方程 : ∂ p t ( x ) ∂ t + ∇ ⋅ ( p t ( x ) u t ( x ) ) = 0 \frac{\partial p_t(x)}{\partial t} +\nabla\cdot\left(p_t(x)u_t(x)\right)=0 ∂t∂pt(x)+∇⋅(pt(x)ut(x))=0

说人话:样本在速度场的作用下不断移动,整个样本群体的分布也会随之变化。连续性方程描述的就是这种变化规律。

  • 论文附录 B 推导了这一关系。借助它,我们可以从速度场的角度描述概率密度如何演化,而不必显式计算流映射的逆映射及其雅可比行列式。
  • 连续性方程建立了速度场与概率密度演化之间的联系。需要注意的是,同一条概率路径可能对应多个不同的速度场,因此我们还需要进一步讨论如何构造合适的概率路径与目标速度。这也正是后续 Flow Matching 训练方法要解决的问题。
2-3-3 代码
  • 于是生成问题被改写成一个回归问题:找一个速度场,使它诱导出的概率路径与目标路径一致
  • 官方代码里这一步的接口就是速度场本身,ModelWrapper 规定了它的形状
python 复制代码
class ModelWrapper(ABC, nn.Module):
    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model

    def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
        return self.model(x=x, t=t, **extras)
  • 输入是位置 x x x 与时刻 t t t,输出就是 u t ( x ) u_t(x) ut(x) 的估计,也就是常微分方程右端那一项
    • x 的形状是 (batch_size, ...),t 是 (batch_size,)。整批样本里的每一条都带自己的时刻,不是整批共用一个 t t t
    • **extras 是留给条件的通道,图像、语言指令、机器人观测都从这里进网络。第三期讲的条件扩散需要额外一路条件输入,Flow Matching 用同样的方式接
    • 库本身既不实现这个方程的积分,也不构造 p t p_t pt。前者交给求解器(见 3-5),后者由调度器提供(见 2-4-4)

说人话:连续性方程说的是"概率质量守恒",速度场是搬运工,概率被搬到哪里、搬得多快,全由它决定

2-4 条件流匹配(Conditional Flow Matching)
  • 上一节我们介绍了速度场与概率分布演化之间的关系。接下来的问题是:如何构造训练目标,让神经网络学会预测合适的速度场?
    • 直接学习边缘概率路径对应的速度场并不容易,因此 Flow Matching 引入了条件流匹配(Conditional Flow Matching,CFM),通过更容易构造的条件概率路径及其对应的速度场来训练模型。
2-4-1 边缘速度场与条件速度场
  • 在介绍具体的训练方法之前,我们需要先区分两个概念:边缘速度场(Marginal Velocity Field) 和条件速度场(Conditional Velocity Field) 。
    • 边缘速度场 u t ( x ) u_t(x) ut(x):描述整体概率分布在时刻 t t t 的运动趋势,它只看当前位置,不管这批样本最终要去哪
    • 条件速度场 u t ( x ∣ x 1 ) u_t(x\mid x_1) ut(x∣x1):给定目标样本 x 1 x_1 x1 之后,对应条件路径的运动速度
  • 用机器人导航打个比方:不指定目标时,需要综合不同目标对应的运动方向,得到的是边缘速度场 ;指定具体目标后,就能根据该目标确定对应的运动方向,这就是条件速度场。
  • 边缘速度场需要对所有可能的目标进行后验加权平均: u t ( x ) = ∫ u t ( x ∣ x 1 ) , p t ( x 1 ∣ x ) , d x 1 u_t(x)=\int u_t(x\mid x_1),p_t(x_1\mid x),dx_1 ut(x)=∫ut(x∣x1),pt(x1∣x),dx1其中, p t ( x 1 ∣ x ) p_t(x_1\mid x) pt(x1∣x) 表示已知当前位置 x x x时,目标为 x 1 x_1 x1 的后验概率分布。计算这个积分通常比较困难,因此难以直接将边缘速度场作为训练监督信号。

也就是说,我们最终需要学习的是边缘速度场,但它通常难以直接计算。因此,CFM 转而利用容易构造的条件速度场作为监督信号训练网络,从而间接学习边缘速度场。

2-4-2 条件流匹配的核心思路
  • CFM 的核心是:不直接计算难以获得的边缘速度场,而是构造容易处理的条件概率路径,并用对应的条件速度作为训练目标。
    • 每次采样一对 ( x 0 , x 1 ) (x_0,x_1) (x0,x1) 和一个时刻 t t t,就能计算路径上的样本及其目标速度,无需求解 ODE 或显式计算边缘概率密度。
  • 那么,使用条件速度训练,能否学到我们真正需要的边缘速度场?在适当条件下,CFM 与 FM 的训练目标具有相同的参数梯度,具体证明见 2-4-6。
  • 接下来,我们先构造一般形式的条件概率路径,再选择具体路径并推导其速度,最后分析条件速度场与边缘速度场的关系。
2-4-3 构造高斯条件路径
  • 我们希望构造一条概率路径,使其从标准高斯噪声分布 π 0 \pi_0 π0 过渡到真实数据分布 q q q: p 0 = π 0 , p 1 = q p_0=\pi_0,\qquad p_1=q p0=π0,p1=q这里的 0 0 0 与 1 1 1 分别表示生成过程的起点与终点

  • 直接构造边缘概率路径比较困难,因此先固定一个目标样本 x 1 x_1 x1,构造对应的条件路径: p 0 ( x ∣ x 1 ) = π 0 ( x ) , p 1 ( x ∣ x 1 ) = δ ( x − x 1 ) p_0(x\mid x_1)=\pi_0(x),\qquad p_1(x\mid x_1)=\delta(x-x_1) p0(x∣x1)=π0(x),p1(x∣x1)=δ(x−x1)其中, δ ( x − x 1 ) \delta(x-x_1) δ(x−x1) 表示狄拉克分布,可以直观地理解为:概率质量全部集中在 x 1 x_1 x1 这一点上,其他位置的概率质量为 0。

  • 再按照真实数据分布 q ( x 1 ) q(x_1) q(x1) 混合所有条件路径,即可得到边缘概率路径: p t ( x ) = ∫ p t ( x ∣ x 1 )   q ( x 1 )   d x 1 p_t(x)=\int p_t(x\mid x_1)\, q(x_1)\, dx_1 pt(x)=∫pt(x∣x1)q(x1)dx1

  • 因此,我们不必直接求解复杂的边缘路径,可以先构造容易处理的条件路径。本节采用高斯条件路径 : p t ( x ∣ x 1 ) = N ( x ∣ α t x 1 ,   σ t 2 I ) p_t(x\mid x_1)=\mathcal{N}\left(x\mid \alpha_t x_1,\, \sigma_t^2 I\right) pt(x∣x1)=N(x∣αtx1,σt2I)也就是把每个时刻的分布取成均值在 α t x 1 \alpha_t x_1 αtx1 上的高斯

  • 从采样角度看,等价于将目标数据与标准高斯噪声按时间相关的系数混合: x t = α t x 1 + σ t x 0 , x 0 ∼ N ( 0 , I ) x_t=\alpha_t x_1+\sigma_t x_0,\qquad x_0\sim\mathcal{N}(0,I) xt=αtx1+σtx0,x0∼N(0,I)其中, α t \alpha_t αt 和 σ t \sigma_t σt 分别控制数据与噪声的贡献,并满足边界条件: α 0 = 0 , α 1 = 1 , σ 0 = 1 , σ 1 = 0 \alpha_0=0,\quad \alpha_1=1,\qquad \sigma_0=1,\quad \sigma_1=0 α0=0,α1=1,σ0=1,σ1=0

  • 这样, t = 0 t=0 t=0 时得到纯噪声, t = 1 t=1 t=1 时得到目标数据。接下来,只需选择具体的系数函数,并推导对应的条件速度场。

2-4-4 构造线性插值路径
2-4-4-1 CondOT(Conditional Optimal Transport,条件最优传输)
  • 上一节给出了条件路径的通用形式: x t = α t x 1 + σ t x 0 x_t=\alpha_t x_1+\sigma_t x_0 xt=αtx1+σtx0其中, x 0 x_0 x0 是标准高斯噪声, x 1 x_1 x1 是真实数据样本。那么,如何选择 α t \alpha_t αt 和 σ t \sigma_t σt?

  • 本期采用一种经典的构造方式:CondOT ,它选了一组最简的线性调度函数 α t = t , σ t = 1 − t \alpha_t=t,\qquad \sigma_t=1-t αt=t,σt=1−t

  • 代入后得到线性插值路径: x t = ( 1 − t ) x 0 + t x 1 \boxed{x_t=(1-t)x_0+tx_1} xt=(1−t)x0+tx1当 t = 0 t=0 t=0 时, x t = x 0 x_t=x_0 xt=x0;当 t = 1 t=1 t=1 时, x t = x 1 x_t=x_1 xt=x1。对于固定的起点和终点,样本会沿两点之间的直线移动。

2-4-4-2 最优传输
  • 最优传输(Optimal Transport,OT)研究如何以最小运输代价,将一个概率分布转移到另一个概率分布。
  • 需要注意:单条路径是直线,不代表所有起点与终点的配对都是全局最优的运输方案。 最优传输性质还取决于样本的配对方式,后续在 2-6 节介绍的 Rectified Flow 会进一步讨论这一问题。
2-4-4-3 代码实现
  • 官方实现中,CondOTScheduler 提供时间调度参数,AffineProbPath 根据这些参数计算路径点与路径速度。
python 复制代码
def __call__(self, t: Tensor) -> SchedulerOutput:
    return SchedulerOutput(
        alpha_t=t,
        sigma_t=1 - t,
        d_alpha_t=torch.ones_like(t),
        d_sigma_t=-torch.ones_like(t),
    )
  • 除了返回 α t = t \alpha_t=t αt=t 和 σ t = 1 − t \sigma_t=1-t σt=1−t,调度器还返回它们对时间的导数: α ˙ t = 1 , σ ˙ t = − 1 \dot\alpha_t=1,\qquad \dot\sigma_t=-1 α˙t=1,σ˙t=−1通用路径逻辑为:
python 复制代码
x_t = sigma_t * x_0 + alpha_t * x_1
dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
  • 第二行是对路径表达式求导: x ˙ t = σ ˙ t x 0 + α ˙ t x 1 \dot x_t=\dot\sigma_t x_0+\dot\alpha_t x_1 x˙t=σ˙tx0+α˙tx1代入 CondOT 的导数,得到: x ˙ t = x 1 − x 0 \boxed{\dot x_t=x_1-x_0} x˙t=x1−x0

  • 也就是说,对于固定的起点和终点,线性路径的速度恒定。

2-4-4-4 从路径速度到条件速度场
  • 路径速度 x ˙ t = x 1 − x 0 \dot x_t=x_1-x_0 x˙t=x1−x0 依赖起点 x 0 x_0 x0,而条件速度场需要用当前位置 x t x_t xt 和目标样本 x 1 x_1 x1 表示。
  • 由线性插值公式,在 t < 1 t<1 t<1 时可得: x 0 = x t − t x 1 1 − t x_0=\frac{x_t-tx_1}{1-t} x0=1−txt−tx1代入路径速度,并将 x t x_t xt 简写为 x x x,得到条件速度场: u t ( x ∣ x 1 ) = x 1 − x 1 − t \boxed{u_t(x\mid x_1)=\frac{x_1-x}{1-t}} ut(x∣x1)=1−tx1−x
  • 因此,给定噪声 x 0 x_0 x0、目标样本 x 1 x_1 x1 和时刻 t t t,就能计算路径点及其对应的条件速度,用作网络的训练目标。
  • 注意,路径速度沿固定路径保持不变,但条件速度场显式依赖 t t t,在 t = 1 t=1 t=1 处不能直接代入计算。两者的关系将在下一节进一步解释。
2-4-5 从条件速度到边缘速度
  • 上一节得到了条件速度场,但真正描述整体概率分布演化的,是边缘速度场。两者之间满足条件期望关系: u t ( x ) = E x 1 − x 0 ∣ x t = x u_t(x)=\mathbb Ex_1-x_0\\mid x_t=x ut(x)=Ex1−x0∣xt=x即在已知当前位置 x x x 的条件下,对所有可能对应的起点与终点所产生的路径速度取平均。由于难以直接计算这个条件期望,CFM 才采用更容易获得的条件速度作为训练监督信号。
  • 对于 CondOT,两个端点的边缘速度可以解析求得。
  • 起点 t = 0 t=0 t=0:此时 x t = x 0 x_t=x_0 xt=x0,而噪声 x 0 x_0 x0 与数据样本 x 1 x_1 x1 独立,因此 u 0 ( x ) = E x 1 − x 0 ∣ x 0 = x = E x 1 − x = μ − x \begin{aligned}u_0(x)&=\mathbb{E}\leftx_1-x_0\\mid x_0=x\\right\\ &=\mathbb{E}\leftx_1\\right-x\\&=\mu-x\end{aligned} u0(x)=Ex1−x0∣x0=x=Ex1−x=μ−x其中 μ \mu μ 是真实数据分布的均值。因此,起点处的边缘速度指向数据均值。
  • 终点 t = 1 t=1 t=1:此时 x t = x 1 x_t=x_1 xt=x1,利用噪声均值为零且与数据独立,可得: E x 1 − x 0 ∣ x 1 = x = x − E x 0 = x \begin{aligned}\mathbb{E}\leftx_1-x_0\\mid x_1=x\\right&=x-\mathbb{E}\leftx_0\\right\\&=x\end{aligned} Ex1−x0∣x1=x=x−Ex0=x这里的 u 1 ( x ) = x u_1(x)=x u1(x)=x 是通过条件期望推导出的解析结果,并不是将 t = 1 t=1 t=1 直接代入条件速度公式。虽然公式中的分母 1 − t 1-t 1−t 趋于零,但沿真实采样路径有: x 1 − x t = ( 1 − t ) ( x 1 − x 0 ) x_1-x_t=(1-t)(x_1-x_0) x1−xt=(1−t)(x1−x0)分子与分母中的 1 − t 1-t 1−t 可以抵消,因此路径速度仍为 x 1 − x 0 x_1-x_0 x1−x0,不会发散。
  • 再看官方代码中的 AffineProbPath.target_to_velocity:
python 复制代码
    def target_to_velocity(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
        scheduler_output = self.scheduler(t)

        alpha_t = scheduler_output.alpha_t
        d_alpha_t = scheduler_output.d_alpha_t
        sigma_t = scheduler_output.sigma_t
        d_sigma_t = scheduler_output.d_sigma_t

        a_t = d_sigma_t / sigma_t
        b_t = (d_alpha_t * sigma_t - d_sigma_t * alpha_t) / sigma_t

        return a_t * x_t + b_t * x_1
  • 它将条件速度写成当前位置与目标样本的线性组合: u t ( x ∣ x 1 ) = σ ˙ t σ t x + α ˙ t σ t − σ ˙ t α t σ t x 1 u_t(x\mid x_1)=\frac{\dot\sigma_t}{\sigma_t}x+\frac{\dot\alpha_t\sigma_t-\dot\sigma_t\alpha_t}{\sigma_t}x_1 ut(x∣x1)=σtσ˙tx+σtα˙tσt−σ˙tαtx1代入 CondOT 的调度参数,就能还原前面的闭式表达式。该实现同样存在 t = 1 t=1 t=1 处分母为零的问题。
  • 本节的关键结论是: 训练时只需通过路径采样获得条件速度作为监督信号,无须显式计算边缘速度场。需要注意的是,网络只在训练采样覆盖的输入上受到监督,面对偏离条件路径的输入,尤其是接近终点的输入时,预测仍可能不稳定。第 4-4 节将进一步观察这一现象。
2-4-6 Flow Matching Loss
  • 理清了条件速度场与边缘速度场的关系后,回到训练目标本身。

  • 如果直接学习边缘速度场 u t ( x ) u_t(x) ut(x),Flow Matching 的损失函数为: L F M ( θ ) = E t ,   p t ( x ) ∥ v θ ( x , t ) − u t ( x ) ∥ 2 \mathcal{L}{\mathrm{FM}}(\theta) = \mathbb{E}{t,\, p_t(x)}\left\| v_\theta(x, t) - u_t(x) \right\|^2 LFM(θ)=Et,pt(x)∥vθ(x,t)−ut(x)∥2

  • 但边缘速度场涉及难以直接计算的条件期望,因此这个目标难以直接用于训练。CFM 则改用条件速度场作为监督信号: L C F M ( θ ) = E t , q ( x 1 ) , p t ( x ∣ x 1 ) ∣ v θ ( x , t ) − u t ( x ∣ x 1 ) ∣ 2 \mathcal L_{\mathrm{CFM}}(\theta) =\mathbb E_{t,q(x_1),p_t(x\mid x_1)} \left\|v_\\theta(x,t)-u_t(x\\mid x_1)\|\^2\\right LCFM(θ)=Et,q(x1),pt(x∣x1)∣vθ(x,t)−ut(x∣x1)∣2其中, x 1 x_1 x1 从数据分布采样, x 0 x_0 x0 从标准高斯分布采样, t t t 从 0 , 1 0,1 0,1 均匀采样。对于 CondOT 路径,路径点与条件速度都有闭式表达,因此训练目标可以直接计算: L C F M ( θ ) = E t , q ( x 1 ) , π 0 ( x 0 ) ∣ v θ ( ( 1 − t ) x 0 + t x 1 , t ) − ( x 1 − x 0 ) ∣ 2 \boxed{ \mathcal L_{\mathrm{CFM}}(\theta) =\mathbb E_{t,q(x_1),\pi_0(x_0)} \left \\left\| v_\\theta((1-t)x_0+tx_1,t) -(x_1-x_0) \\right\|\^2 \\right } LCFM(θ)=Et,q(x1),π0(x0)∣vθ((1−t)x0+tx1,t)−(x1−x0)∣2这就是代码中通过均方误差(Mean Squared Error,MSE)训练速度预测网络的核心逻辑。

  • 那么,使用条件速度作为监督信号,能否让网络学到我们真正需要的边缘速度场?论文定理 2 给出了关键结论:在相应条件下,FM 与 CFM 的目标函数对模型参数具有相同的梯度: ∇ θ L F M ( θ ) = ∇ θ L C F M ( θ ) \boxed{ \nabla_\theta\mathcal L_{\mathrm{FM}}(\theta) =\nabla_\theta\mathcal L_{\mathrm{CFM}}(\theta) } ∇θLFM(θ)=∇θLCFM(θ)其核心是条件期望的平方误差分解。对固定的 x , t x,t x,t,有: ∥ v θ ( x ) − E u t ( x ∣ x 1 ) ∥ 2 = E ∥ v θ ( x ) − u t ( x ∣ x 1 ) ∥ 2 − V a r u t ( x ∣ x 1 ) \left\|v_\theta(x) - \mathbb{E}\leftu_t(x\|x_1)\\right\right\|^2 = \mathbb{E}\left\|v_\theta(x) - u_t(x|x_1)\right\|^2 - \mathrm{Var}\leftu_t(x\|x_1)\\right ∥vθ(x)−Eut(x∣x1)∥2=E∥vθ(x)−ut(x∣x1)∥2−Varut(x∣x1)

  • 这里,边缘速度场是给定当前位置后,条件速度的后验均值。等式右侧的第一项对应条件速度的平方误差,第二项是条件速度的后验方差。由于该方差不依赖模型参数 (\theta),对参数求梯度时这一项便会消失,从而得到两个目标函数的梯度等价性。

说人话:训练时,网络拟合的是每条条件路径对应的速度;而在给定当前位置后,这些条件速度的平均值就是边缘速度。梯度等价意味着,用条件速度训练并不会改变优化方向,因此我们可以用容易计算的条件目标来训练模型,而不必直接计算难以获得的边缘速度场。

  • 落到官方代码中,训练目标就是以下几行:
python 复制代码
    # sample probability path
    path_sample = path.sample(t=t, x_0=x_0, x_1=x_1)

    # flow matching l2 loss
    loss = torch.pow( vf(path_sample.x_t,path_sample.t) - path_sample.dx_t, 2).mean()
  • path.sample 返回路径上的样本 x_t 及对应的目标速度 dx_t,网络 vf 根据 x_t 和 t 预测速度,再通过均方误差计算损失。整个训练过程不需要显式计算边缘速度场,也不需要在每次训练时求解 ODE。
2-5 训练与推理流程
  • 有了上面的目标函数,训练循环就变得非常简洁。推理阶段也不复杂:从噪声出发,通过数值积分不断更新样本的位置。
  • 本节结合官方示例,分别介绍训练与推理的代码流程。
2-5-1 训练:随机采样时间并拟合速度

#mermaid-svg-wNrctbo5HHXBCEqJ{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-wNrctbo5HHXBCEqJ .error-icon{fill:#552222;}#mermaid-svg-wNrctbo5HHXBCEqJ .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-wNrctbo5HHXBCEqJ .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-wNrctbo5HHXBCEqJ .marker{fill:#333333;stroke:#333333;}#mermaid-svg-wNrctbo5HHXBCEqJ .marker.cross{stroke:#333333;}#mermaid-svg-wNrctbo5HHXBCEqJ svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-wNrctbo5HHXBCEqJ p{margin:0;}#mermaid-svg-wNrctbo5HHXBCEqJ .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster-label text{fill:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster-label span{color:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster-label span p{background-color:transparent;}#mermaid-svg-wNrctbo5HHXBCEqJ .label text,#mermaid-svg-wNrctbo5HHXBCEqJ span{fill:#333;color:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ .node rect,#mermaid-svg-wNrctbo5HHXBCEqJ .node circle,#mermaid-svg-wNrctbo5HHXBCEqJ .node ellipse,#mermaid-svg-wNrctbo5HHXBCEqJ .node polygon,#mermaid-svg-wNrctbo5HHXBCEqJ .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-wNrctbo5HHXBCEqJ .rough-node .label text,#mermaid-svg-wNrctbo5HHXBCEqJ .node .label text,#mermaid-svg-wNrctbo5HHXBCEqJ .image-shape .label,#mermaid-svg-wNrctbo5HHXBCEqJ .icon-shape .label{text-anchor:middle;}#mermaid-svg-wNrctbo5HHXBCEqJ .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-wNrctbo5HHXBCEqJ .rough-node .label,#mermaid-svg-wNrctbo5HHXBCEqJ .node .label,#mermaid-svg-wNrctbo5HHXBCEqJ .image-shape .label,#mermaid-svg-wNrctbo5HHXBCEqJ .icon-shape .label{text-align:center;}#mermaid-svg-wNrctbo5HHXBCEqJ .node.clickable{cursor:pointer;}#mermaid-svg-wNrctbo5HHXBCEqJ .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-wNrctbo5HHXBCEqJ .arrowheadPath{fill:#333333;}#mermaid-svg-wNrctbo5HHXBCEqJ .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-wNrctbo5HHXBCEqJ .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-wNrctbo5HHXBCEqJ .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-wNrctbo5HHXBCEqJ .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-wNrctbo5HHXBCEqJ .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-wNrctbo5HHXBCEqJ .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster text{fill:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ .cluster span{color:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-wNrctbo5HHXBCEqJ .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-wNrctbo5HHXBCEqJ rect.text{fill:none;stroke-width:0;}#mermaid-svg-wNrctbo5HHXBCEqJ .icon-shape,#mermaid-svg-wNrctbo5HHXBCEqJ .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-wNrctbo5HHXBCEqJ .icon-shape p,#mermaid-svg-wNrctbo5HHXBCEqJ .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-wNrctbo5HHXBCEqJ .icon-shape .label rect,#mermaid-svg-wNrctbo5HHXBCEqJ .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-wNrctbo5HHXBCEqJ .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-wNrctbo5HHXBCEqJ .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-wNrctbo5HHXBCEqJ :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} ① 抽取样本对

噪声 x₀ 与真实数据 x₁
② 随机采样时刻

t ~ U(0, 1)
③ 构造路径点与目标速度

xₜ = (1-t)x₀ + tx₁

uₜ = x₁ - x₀
④ 平方误差回归

L = ||vθ(xₜ,t) - uₜ||²

  • 代码注释中的 user's responsibility 强调的是:数据采样与时间采样由调用方负责,路径类负责根据输入计算路径点和目标速度。
    • x 1 x_1 x1 从数据分布采样,示例中的 inf_train_gen 使用双月形数据生成器; x 0 x_0 x0 从标准高斯分布采样。
    • t t t 从 0 , 1 0,1 0,1均匀采样。实际应用中也可以设计其他时间采样分布,本例采用最简单的均匀采样。
    • x 0 x_0 x0 与 x 1 x_1 x1 独立采样,构成随机配对。这样的配对不一定符合最优传输方案,生成的直线路径也可能相互交叉。第 2-6 节将介绍 Rectified Flow 如何通过重新配对等方式进一步改善轨迹。
  • 核心训练循环如下:
python 复制代码
# instantiate an affine path object
path = AffineProbPath(scheduler=CondOTScheduler())

# init optimizer
optim = torch.optim.Adam(vf.parameters(), lr=lr)

# train
for i in range(iterations):
    optim.zero_grad()

    x_1 = inf_train_gen(batch_size=batch_size, device=device) # sample data
    x_0 = torch.randn_like(x_1).to(device)
    t = torch.rand(x_1.shape[0]).to(device)

    # sample probability path
    path_sample = path.sample(t=t, x_0=x_0, x_1=x_1)

    # flow matching l2 loss
    loss = torch.pow( vf(path_sample.x_t,path_sample.t) - path_sample.dx_t, 2).mean()

    loss.backward()
    optim.step()
  • 这段代码有几个值得关注的地方:

    1. 训练时无需数值求解 ODE。 每次迭代只需采样路径点、计算目标速度,再通过损失函数更新网络参数,也不需要显式计算边缘概率密度。
    2. 目标速度可以直接计算。 对于 CondOT,path_sample.dx_t 对应 x 1 − x 0 x_1-x_0 x1−x0,不依赖网络预测结果。
    3. 同一批次可以覆盖不同的时间位置。 t t t 是逐样本采样的,因此网络能同时学习路径上不同位置的速度。
    4. 网络输入不包含起点与终点的配对信息。 网络接收的是 ( x t , t ) (x_t,t) (xt,t),而不是 ( x 0 , x 1 ) (x_0,x_1) (x0,x1)。训练时,配对信息用于构造监督信号;推理时,网络只需根据当前样本与时间预测速度。
  • 这也是模型能够在推理阶段独立运行的关键:它学到的是速度场,而不是记住训练时的样本配对。

2-5-2 推理:从噪声积分到目标样本
  • 训练完成后,我们不再拥有真实目标样本 x 1 x_1 x1,而是使用训练好的速度网络 v θ v_\theta vθ,从初始噪声 x 0 ∼ N ( 0 , I ) x_0\sim\mathcal N(0,I) x0∼N(0,I) 出发,求解常微分方程: d x t d t = v θ ( x t , t ) \frac{dx_t}{dt}=v_\theta(x_t,t) dtdxt=vθ(xt,t)积分至 t = 1 t=1 t=1,得到生成样本: x 1 = x 0 + ∫ 0 1 v θ ( x t , t ) , d t x_1=x_0+\int_0^1v_\theta(x_t,t),dt x1=x0+∫01vθ(xt,t),dt
  • 官方示例将积分过程交给 ODESolver:
python 复制代码
# step size for ode solver
step_size = 0.05

batch_size = 50000  # batch size
T = torch.linspace(0,1,10)  # sample times
T = T.to(device=device)

x_init = torch.randn((batch_size, 2), dtype=torch.float32, device=device)
solver = ODESolver(velocity_model=wrapped_vf)  # create an ODESolver class
sol = solver.sample(time_grid=T, x_init=x_init, method='midpoint', step_size=step_size, return_intermediates=True)  # sample from the model
  • 其中:
    • wrapped_vf 是前面介绍过的 ModelWrapper,负责统一速度网络的调用接口。
    • time_grid 指定求解器使用的时间区间及输出时刻;step_size 则控制内部积分步长。两者作用不同,具体的离散化方式取决于求解器的实现。
    • method="midpoint" 使用中点法进行积分。作为二阶方法,它通常需要每一步进行两次速度场评估,因此积分步数与网络求值次数(NFE,Number of Function Evaluations)并不相同。
  • 为了让后续实验中的计算成本更容易比较,本期实验改用一阶 Euler 方法,并让时间网格直接决定积分步数,不额外指定 step_size。在这一设置下,每个积分步对应一次网络求值,因此步数与 NFE 一致,具体见第 3-5 节和第 4-5 节。
  • 在固定模型与数值求解设置下,Flow Matching 的采样过程是确定性的:
    • 给定相同的初始噪声,就会得到相同的生成结果。 生成多样性主要来自初始噪声的采样,而不是积分过程中额外注入随机噪声。
  • 这一特性对机器人任务很有价值:相同输入下的生成过程更容易复现,有利于调试与评估。不过,如果初始噪声采样不够充分,生成结果的多样性也可能受到影响。第 4-5 节将通过不同的 NFE 设置观察采样效果与计算成本之间的关系。
2-6 Rectified Flow:让生成路径更直接
  • 先说清本节的范围:本文实验用的是随机配对的 CondOT,没有做重配对训练,所以严格来说它只对应 Rectified Flow 的初始路径设置,而不是完成了 rectification 的模型。这一区分决定了后文哪些读数可以拿来比较。
  • 2-5 提到过,由于 x 0 x_0 x0 与 x 1 x_1 x1 独立采样,不同样本的直线路径可能在中途相交。但真正起作用的机制比"交叉"更一般:在同一时空位置上,不同条件路径给出的速度监督会彼此冲突 。
    • 平方误差的最优解是这些条件速度的均值,因此模型学到的是平均方向,而不是其中任何一条路径的方向。沿着这个平均速度场积分,最终得到的轨迹就不一定是直线。
    • 轨迹交叉是这种冲突最直观的例子,却不是唯一来源。即使几何上看不出交叉,随机配对也可能让边缘流线变弯。
  • Rectified Flow(Liu 等人,2022)研究的正是如何改善这种传输路径。它的基本训练目标是: min ⁡ v ∫ 0 1 E ∣ ( x 1 − x 0 ) − v ( x t , t ) ∣ 2 d t , x t = ( 1 − t ) x 0 + t x 1 \min_v \int_0^1 \mathbb{E}\left\\left\|(x_1-x_0)-v(x_t,t)\\right\|\^2\\rightdt, \qquad x_t=(1-t)x_0+tx_1 vmin∫01E∣(x1−x0)−v(xt,t)∣2dt,xt=(1−t)x0+tx1
  • 可以看到,它与前面介绍的 CondOT 条件流匹配损失形式相同,使用的也是线性插值路径。模型根据随机配对 ( x 0 , x 1 ) (x_0,x_1) (x0,x1) 生成训练样本,并回归两者之间的位移 x 1 − x 0 x_1-x_0 x1−x0。
  • 在平方误差目标下,最优速度场为条件期望: v X ( x , t ) = E x 1 − x 0 ∣ x t = x v^X(x,t)=\mathbb{E}x_1-x_0\\mid x_t=x vX(x,t)=Ex1−x0∣xt=x
  • 这与前面介绍的 CondOT 边缘速度场是一致的。因此,Rectified Flow 的关键并不只是定义一个新的训练损失,而是通过重新配对进一步改善传输路径。
  • 具体来说,记第 k k k 轮的配对为 Z k = ( Z 0 k , Z 1 k ) Z^k=(Z_0^k,Z_1^k) Zk=(Z0k,Z1k),每轮主要包含三个步骤:
    • 使用当前速度场 v k v_k vk,从 t = 0 t=0 t=0积分到 t = 1 t=1 t=1,将每个起点 Z 0 k Z_0^k Z0k 映射到模型生成的终点。
    • 将原起点与新的终点组成配对,得到下一轮训练数据。
    • 使用新配对重新训练速度场 v k + 1 v_{k+1} vk+1。这一过程称为 rectification,可以重复多轮。
  • 为什么重新配对有助于让路径变直?第一轮使用随机配对,不同样本的目标方向可能相互冲突。经过一次 ODE 积分后,新的终点由当前速度场生成,起点与终点之间的对应关系也随之改变。重新训练后,模型有机会学习到交叉更少、更加直接的传输路径。需要注意的是,路径并不保证每轮都严格变直。
  • 论文还给出了关于路径直线程度的定量界。其定理 3.7 可写为: min ⁡ k S ( Z k ) ≤ E ∣ x 1 − x 0 ∣ 2 K \min_{k} S(Z^k)\leq \frac{\mathbb{E}\left\|x_1-x_0\|\^2\\right}{K} kminS(Zk)≤KE∣x1−x0∣2其中, K K K 表示考虑的轮数, S S S 用来衡量轨迹偏离直线运动的程度: S ( Z ) = ∫ 0 1 E ∣ ( Z 1 − Z 0 ) − Z ˙ t ∣ 2 d t S(Z)=\int_0^1 \mathbb{E}\left \\left\|(Z_1-Z_0)-\\dot Z_t\\right\|\^2 \\rightdt S(Z)=∫01E (Z1−Z0)−Z˙t 2dt
  • 直观来说, Z 1 − Z 0 Z_1-Z_0 Z1−Z0 是起点到终点的平均位移方向,而 Z ˙ t \dot Z_t Z˙t 是轨迹在时刻 t t t 的实际速度。两者越接近, S S S 就越小;若轨迹以恒定速度沿直线运动,则 S = 0 S=0 S=0。
    • 这个界说明,在定理所需的假设下, K K K 轮中至少存在一轮,其路径直线程度指标不超过上述上界。因此,增加重配对轮数有机会改善路径,但不能简单理解为每轮都会让路径更直,也不能据此直接保证采样步数按 (1/K) 下降。
  • 论文也提醒,rectification 并非轮数越多越好。一方面,每轮都需要重新生成配对并训练模型,计算成本会随轮数增加;另一方面,反复重配对会改变学习到的传输耦合及生成分布,过多迭代可能损害生成质量。因此,更适合将其视为改善传输路径的精修过程,而不是无限重复的优化步骤。
  • 回到本节开头那句范围说明:本文实验停在初始路径设置这一步,第 4 章量到的都是 CondOT 的表现,不包含 rectification 带来的改善。

3 Flow Matching 网络架构与实现

3-1 整体架构
  • 官方代码将 Flow Matching 拆分为几个职责清晰、耦合较弱的模块:
    • 路径与调度器 :决定概率路径的形式,提供 α t \alpha_t αt、 σ t \sigma_t σt 及其导数。
    • 速度场模型 :接收当前位置 x x x 与时间 t t t,输出预测速度。
    • 求解器:从初始噪声出发,对速度场对应的 ODE 进行数值积分。
  • 从接口设计上看,这种解耦十分直观:
    • 路径主要用于训练,负责生成路径点和监督目标;推理时通常不需要它。
    • 模型只接收 ( x , t ) (x,t) (x,t),不需要知道训练时使用了什么路径,也不需要知道求解器如何积分。
    • 求解器只需要能够调用速度场,例如 model(x=x, t=t),不关心模型内部使用的是 MLP(Multi-Layer Perceptron,多层感知机)还是其他网络结构。
  • 因此,在接口兼容的前提下,更换路径、网络或求解器,通常不需要同步修改其他模块。各组件的分工与使用时机如下。
组件 官方类或模块 职责 训练时 推理时
路径 AffineProbPath、CondOTProbPath 生成路径点 x t x_t xt 与条件速度 使用 通常不使用
调度器 CondOTScheduler、VPScheduler 提供 α t \alpha_t αt、 σ t \sigma_t σt 及其导数 使用 通常不使用
模型 MLP、ModelWrapper 学习速度场 使用 使用
求解器 ODESolver 对速度场进行数值积分 不使用 使用
损失 训练循环中的平方误差 衡量预测速度与监督目标的差异 使用 不使用
  • 损失这一行也值得注意:在本文所使用的连续 Flow Matching 示例中,训练损失直接写在训练循环里,并没有专门的损失类。库中的 MixturePathGeneralizedKL 则面向离散数据的相关目标。
  • 本期实验沿用这套框架,只调整网络宽度,将官方示例的 128 提高到 256,其余路径、调度器、求解器和损失设计保持不变。这样,后续实验就能在相对一致的实现条件下,考察采样步数与生成质量之间的关系。
3-2 时间嵌入与条件输入
  • 网络需要接收连续时间变量 t t t。常见做法是将时间编码为高维正弦嵌入,但官方示例更简单:直接将时间标量拼接到输入坐标后。
python 复制代码
def forward(self, x: Tensor, t: Tensor) -> Tensor:
    sz = x.size()
    x = x.reshape(-1, self.input_dim)
    t = t.reshape(-1, self.time_dim).float()

    t = t.reshape(-1, 1).expand(x.shape[0], 1)
    h = torch.cat([x, t], dim=1)
    output = self.main(h)

    return output.reshape(*sz)
  • 代码先将 x 展平为二维张量,将 t 转为浮点数并扩展为每个样本对应的标量,最后将输出恢复为输入的原始形状。
  • 为什么不使用正弦时间嵌入? 这是一个二维玩具任务,直接拼接时间即可。虽然 CondOT 单条条件路径的速度 (x_1-x_0) 与时间无关,但模型学习的是边缘速度场,不同时间的边缘分布仍然不同,因此 (t) 依然有用。对于图像生成或机器人动作生成等任务,时间嵌入可能更有帮助,但并非高维任务的硬性要求。
  • 此外,接口 forward(self, x, t, **extras) 还支持传入额外条件。求解器会将这些参数转发给速度场模型,由模型决定如何融合条件信息,而不会改变 ODE 数值积分的基本逻辑。本期实验是无条件生成,因此 **extras 为空;这一设计也为后续扩展条件生成预留了接口。
3-3 速度预测网络
  • 网络主体是一个使用 Swish 激活函数的 MLP,官方示例的定义如下:
python 复制代码
# Activation class
class Swish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: Tensor) -> Tensor:
        return torch.sigmoid(x) * x


# Model class
class MLP(nn.Module):
    def __init__(
        self,
        input_dim: int = 2,
        time_dim: int = 1,
        hidden_dim: int = 128,
    ):
        super().__init__()

        self.input_dim = input_dim
        self.time_dim = time_dim
        self.hidden_dim = hidden_dim

        self.main = nn.Sequential(
            nn.Linear(input_dim + time_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, input_dim),
        )
  • 从结构上看,该 MLP 并不复杂:输入经过隐藏层的线性变换与激活,最终输出与数据维度一致的速度向量。本期实验中,输入为二维坐标加一维时间,共 3 维;输出为二维速度。
    • 激活函数使用 Swish :其表达式为 x σ ( x ) x\sigma(x) xσ(x),相比 ReLU 在零点更平滑,但不代表实际预测精度一定更高。
    • 没有使用归一化层:对于低维玩具任务,保持结构简单即可,是否需要归一化仍需结合实验判断。
  • 本次实验主要调整网络宽度 hidden_dim。设隐藏维度为 (h),按照该网络结构,参数总量为: P = 3 h 2 + 9 h + 2 P=3h^2+9h+2 P=3h2+9h+2
  • 不同宽度对应的参数量如下:
hidden_dim 参数量
128(官方默认) 50,306
256(本期实验) 198,914
512(官方 Notebook) 791,042
  • 本期取 h = 256 h=256 h=256,在模型容量与实验预算之间折中。若网络容量不足,生成质量可能提前受限,影响对训练步数的比较;但该配置是否足够,仍需通过实验验证。
  • 相比文本或图像任务中可能使用的卷积网络、Transformer 等架构,这个 MLP 简单得多,但核心任务相同:根据当前状态和时间预测速度
3-4 训练代码解析
  • 本节在官方训练循环的基础上,进一步讨论几个关键设计点。
    • 训练配置 :官方示例使用批量大小 4096、训练 20,000 步,并采用学习率为 10 − 3 10^{-3} 10−3 的 Adam 优化器。每一步都会重新采样噪声、数据和时间,因此用训练步数描述预算比 epoch 更直观。
    • 验证与梯度裁剪:官方示例没有设置验证集、早停或梯度裁剪,属于玩具任务中的简化设计,并不意味着在线采样不会过拟合,或梯度裁剪永远没有必要。后续将通过生成样本和质量指标观察训练效果。
    • 训练目标的构造 :Flow Matching 无需人工标注速度。每一步重新采样 x 0 x_0 x0、 x 1 x_1 x1 和 t t t,由路径模块计算中间状态 x t x_t xt 与条件速度 x ˙ t \dot{x}_t x˙t,分别作为网络输入和监督目标。
    • 训练与推理的区别 :训练时,path.sample 返回路径点与条件速度,用于计算回归损失;推理时,则从初始噪声出发,反复调用速度场求解 ODE。
  • 因此,训练开销主要取决于批量大小、网络规模和训练步数,而不取决于推理时 ODE 求解器使用多少个积分步。两者职责清晰:训练学习速度场,推理利用速度场生成样本。
3-5 ODE 数值积分与采样代码
  • 训练完成后,生成样本的过程就是求解速度场对应的常微分方程。官方将这项工作交给 ODESolver,其核心方法如下:
python 复制代码
def sample(
    self,
    x_init: Tensor,
    step_size: Optional[float],
    method: str = "euler",
    atol: float = 1e-5,
    rtol: float = 1e-5,
    time_grid: Tensor = torch.tensor([0.0, 1.0]),
    return_intermediates: bool = False,
    enable_grad: bool = False,
    **model_extras,
) -> Union[Tensor, Sequence[Tensor]]:
    time_grid = time_grid.to(x_init.device)

    def ode_func(t, x):
        return self.velocity_model(x=x, t=t, **model_extras)

    ode_opts = {"step_size": step_size} if step_size is not None else {}

    with torch.set_grad_enabled(enable_grad):
        sol = odeint(
            ode_func,
            x_init,
            time_grid,
            method=method,
            options=ode_opts,
            atol=atol,
            rtol=rtol,
        )

    if return_intermediates:
        return sol
    else:
        return sol[-1]
  • ode_func(t, x) 调用速度场,返回当前状态的变化率;真正的积分由 torchdiffeq.odeint 完成。官方 ODESolver 则封装了模型调用、积分参数和中间状态返回等功能。
  • 几个关键参数如下:
    • time_grid:指定积分时间范围及输出时刻。
    • step_size:控制积分器内部步长,与输出时间网格并不完全等价;固定步长方法和自适应方法的处理方式也不同。
    • atol、rtol:自适应积分器的绝对与相对误差容差,通常不控制普通固定步长方法的步长。
    • enable_grad=False:默认不记录采样过程的梯度,减少不必要的计算图开销。
    • **model_extras:将额外条件传递给速度场模型。
    • return_intermediates:控制是否返回整条积分轨迹;关闭时只返回最终状态,开启时可用于可视化生成过程。
  • method 指定使用哪一种积分方法。几种常见选择如下。
method 阶数 典型网络求值次数 说明
euler 一阶 每步 1 次 实现简单,便于分析步数与误差的关系
midpoint 二阶 每步 2 次 利用中点处的速度提高积分精度
heun3 三阶 每步 3 次 更高阶的固定步长方法
dopri5 自适应阶数 不固定 根据误差估计动态调整内部步长
  • 固定步长方法的网络调用量可近似按"积分步数 × 每步求值次数"估算;自适应方法则应统计实际的函数求值次数。
  • 本期统一使用 Euler 法。对于积分区间 0 , 1 0,1 0,1,若划分为 N N N 个等长区间,则步长为 h = 1 / N h=1/N h=1/N,对应 N N N 次速度场求值。选择 Euler 是为了让积分步数与网络调用量的关系更直观,方便后续比较采样成本与生成质量。
3-6 采样步数与生成质量
  • 对于足够光滑的速度场,Euler 法的数值误差通常具有以下阶数:
    • 局部截断误差 :单步误差为 O ( h 2 ) O(h^2) O(h2)。
    • 全局误差 :在常见条件下为 O ( h ) O(h) O(h)。对于区间 0 , 1 0,1 0,1,使用 N = 1 / h N=1/h N=1/h 个等长积分区间时,全局误差约为 O ( 1 / N ) O(1/N) O(1/N)。
  • 但这描述的是数值积分误差,并不意味着生成质量会严格按照 1 / N 1/N 1/N 改善。步数较少时,积分误差可能是主要瓶颈;步数增加后,模型拟合误差和质量指标的统计波动可能变得更加显著。
    • 最终生成质量还取决于速度场的拟合难度 和目标分布的复杂程度。例如,CondOT 的条件路径虽然是直线,但边缘速度场不一定恒定,因此不能仅凭路径直线就断言模型容易拟合。相较于二维双月任务,图像生成和机器人动作生成通常还面临更复杂的模型与数据需求。
    • 与此同时,采样成本随积分步数增加。对于 Euler 法,每步调用一次速度场,因此单个样本需要 N N N 次网络求值,实际耗时则取决于模型、硬件和实现方式。
  • 因此,4-5 节将围绕一个核心问题展开:生成质量达到可接受水平时,最少需要多少次网络求值? 我们将在二维双月数据集上扫描不同采样步数,同时记录生成质量与单样本耗时,分析数值精度与推理成本之间的权衡。对于机器人动作生成,推理延迟还必须满足实际控制系统的时间预算,但这一点需要在目标硬件上单独测试。

4 实验:从噪声生成数据

4-1 二维合成数据集
  • 本章的三个实验都跑在双月形数据上。选它的理由是它能暴露一类具体的失败模式
    • 两个交错的弯月,每个月的两端收尖,分布整体是多峰的
    • 用一个单步回归网络去拟合它,网络会输出两个月的平均,也就是夹在中间的一团
    • 这正是第三期讲 Diffusion Policy 时那个多模态问题的最简版本,也是我们要生成模型去解决的原型问题
  • 数据由 scikit-learn 的 make_moons 生成,再叠一层高斯噪声,把离散的曲线铺成有宽度的带
python 复制代码
def two_moons(
    n: int,
    noise: float = 0.05,
    seed: int = 0,
    device: str | torch.device = "cpu",
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """抽 `n` 个双月形样本,返回 (n, 2)。

    同一个 `seed` 必得同一批点;换 `seed` 就是"同一分布的另一批样本",
    于是评测时可以拿它当留出集用,不必真的划分训练/测试。
    """
    pts, _ = make_moons(n_samples=n, noise=noise, random_state=seed)
    return torch.as_tensor(np.asarray(pts), dtype=dtype, device=device)
  • 两个细节值得说明
    • 不做标准化。两组月的坐标本来就落在 − 1.2 , 2.2 × − 0.7 , 1.3 -1.2, 2.2 \times -0.7, 1.3 −1.2,2.2×−0.7,1.3,与源分布 N ( 0 , I ) \mathcal{N}(0, I) N(0,I) 是同一个量级;标准化反而会把月亮压扁,破坏它可辨认的形状
    • 用 seed 换一批点当留出集。同一个 seed 必得同一批点,所以"训练用的数据"与"评测用的数据"只是两个不同的 seed,不需要真的划分
  • 训练时数据先一次抽成一个池子,每个训练步从里面随机取一个小批量,这就是" x 1 x_1 x1 来自数据集"的字面写法
  • 下图左边是留出集,右边是训练好的模型用 32 步 Euler 采样出来的样本
  • 两张图看上去是同一批点。看图只能给出一个印象,4-3 会给这件事配一个标量
4-2 训练过程与损失曲线
  • 训练配置如下,除了网络宽度,其余都沿用官方例子的取值
项目 取值
网络 4 层 Swish MLP,隐层宽度 256
参数量 198,914
优化器 Adam
学习率 10 − 3 10^{-3} 10−3
批量 4096
训练步数 120,000
数据池 100,000 点,噪声标准差 0.05
随机种子 0
  • 训练循环就是 3-4 里那五行,核心部分没有别的东西
python 复制代码
    for step in range(1, steps + 1):
        idx = torch.randint(0, pool_size, (batch_size,), device=device, generator=g)
        x_1 = pool[idx]
        x_0 = torch.randn(batch_size, 2, device=device, generator=g)
        t = torch.rand(batch_size, device=device, generator=g)

        sample = path.sample(x_0=x_0, x_1=x_1, t=t)
        loss = torch.nn.functional.mse_loss(model(sample.x_t, t), sample.dx_t)

        optim.zero_grad(set_to_none=True)
        loss.backward()
        optim.step()
  • 逐行对到 2-5-1 的四件事上:取 x 1 x_1 x1、抽 x 0 x_0 x0 与 t t t、用 path.sample 拼出路径点与目标速度、一行 MSE 回归。梯度经过的只有 model,path 不参与求导

  • 损失曲线如下,纵轴用线性刻度,每个记录点是最近 100 步的平均

  • 曲线形状很干净:前一千步内下降约两成,此后转入缓慢爬行,两万步之后基本是一条水平线

  • 全程没有尖峰,也没有震荡。这就是 3-4 说"不需要梯度裁剪"的依据

  • 但这条曲线里藏着一个陷阱,值得单独说清楚

4-2-1 损失平了不等于质量平了
  • 损失是这个任务里最容易被误读的一个量。它的目标是条件速度 x 1 − x 0 x_1 - x_0 x1−x0,而 x 0 x_0 x0 是标准高斯、 x 1 x_1 x1 是双月形,两者的差本身就是一个随机量,均方误差不可能降到零
  • 更关键的是它每一批都在换数据。每个训练步的配对都是现抽的,所以损失里有一部分是采样噪声,不是模型误差
  • 把训练预算扫一遍就能看出这件事。四档预算训出来的模型,损失几乎一模一样,生成质量却差着一截
训练步数 末段损失 能量距离( N = 128 N = 128 N=128)
20,000 0.9404 1.12 × 10 − 3 1.12 \times 10^{-3} 1.12×10−3
60,000 0.9386 8.2 × 10 − 4 8.2 \times 10^{-4} 8.2×10−4
120,000 0.9392 8.5 × 10 − 4 8.5 \times 10^{-4} 8.5×10−4
200,000 0.9403 7.4 × 10 − 4 7.4 \times 10^{-4} 7.4×10−4
  • 算力翻了十倍,末段损失几乎没动,能量距离却好了约一半
  • 这说明两件事
    • 损失早就到平台了,但它反映的是"网络在随机配对上的平均回归误差",不是"生成出来的分布对不对"
    • 质量也在六万步之后进入平台,再往上加预算只是在种子噪声里晃,所以最终取 120,000 步,是在平台之上留了一倍余量
  • 这条教训是通用的:损失曲线平了,不等于样本分布停止变好。判收敛要看生成质量本身,不能只看损失

说人话:损失量的是"平均猜得准不准",质量量的是"撒出来的点像不像",这在生成模型上不是同一件事

4-3 训练并生成双月形数据
  • 判据用能量距离(ED,Energy Distance)。两个样本集之间的定义为

D 2 = 2   E ∥ x − y ∥ − E ∥ x − x ′ ∥ − E ∥ y − y ′ ∥ D^2 = 2\,\mathbb{E}\|x - y\| - \mathbb{E}\|x - x'\| - \mathbb{E}\|y - y'\| D2=2E∥x−y∥−E∥x−x′∥−E∥y−y′∥

  • 它是 MMD(Maximum Mean Discrepancy,最大均值差异)取特定核时的特例,非负,且两个分布相同时取 0
  • 选它而不选带高斯核的 MMD,理由是不用挑核带宽。带宽一调,不同步数之间的比较就依赖于那个手调参数,而本期实验的自变量只有步数一个
  • 实现上有一个约定要说清楚:自距离项按 V-统计量算,也就是把 i = j i = j i=j 的对角线也算进去
python 复制代码
def _mean_pairwise(a: torch.Tensor, b: torch.Tensor, block: int = 1024) -> float:
    """E‖a - b‖,分块算免得一次性铺开整张 (n, m) 距离矩阵。"""
    total, count = 0.0, 0
    for i in range(0, a.shape[0], block):
        chunk = a[i : i + block]
        total += torch.cdist(chunk, b).sum().item()
        count += chunk.shape[0] * b.shape[0]
    return total / count

def energy_distance(x: torch.Tensor, y: torch.Tensor) -> float:
    """样本集 `x` 与 `y` 之间的能量距离(越小越接近,同分布时为 0)。"""
    if x.shape[0] < 2 or y.shape[0] < 2:
        raise ValueError("两边的样本数都要 >= 2")
    return (
        2.0 * _mean_pairwise(x, y)
        - _mean_pairwise(x, x)
        - _mean_pairwise(y, y)
    )
  • 这么算换来两个性质:结果恒非负,原因是它恰好等于两个经验分布之间的能量距离,而 ∥ ⋅ ∥ \|\cdot\| ∥⋅∥ 是条件负定核;且 x x x 与 y y y 是同一批点时精确等于 0,代码跑出来确实是 0.0
  • 代价是有 O ( 1 / n ) O(1/n) O(1/n) 的上偏 。当 x x x 与 y y y 是同一分布的两批独立 样本时,估计的期望不是 0,而是 2 n   E ∥ x − x ′ ∥ \frac{2}{n}\,\mathbb{E}\|x - x'\| n2E∥x−x′∥
  • 但"同一分布"的读数并不为零,地板要单独量:取若干对独立批次估它,读数比上面那个理论上偏小一个数量级以上
    • 原因在数据本身。make_moons 的两批点共用同一条确定性的骨架,换种子只换上面那层噪声,真实的自距离项远小于独立同分布假设下的量级;而模型要复现的也正是这条骨架,用同一套抽法量出来的地板对它是同一个口径
  • 结果:留出样本上做 32 步 Euler 采样,能量距离离地板还有一个多数量级
  • 但这是一次抽样的结果,同一批权重换采样种子重来能差出好几倍,所以单次读数偏乐观
    • 后文所有跨步数的比较因此统一用多种子均值;一个读数也只在乘以几倍的差别上才有意义,百分之几十的高低读不出来
  • 距离数之外,再用最近邻距离核对点的疏密。生成点到真实点的最近邻距离中位数,与真实点之间的同一个量几乎相同,方向上略紧一点
    • 取真实数据自身最近邻距离的九十分位数当阈值,生成点里超过它的比例低于随机情况下的基准;一阶矩与二阶矩也指向同一结论,均值几乎重合,标准差都略小一点
    • 三条读数方向一致:生成分布比真实数据略紧,不是更散

说人话:这里的情况和通常的担心相反,模型不是把分布铺得太开,而是略微收紧了一点

4-4 速度场与生成轨迹
  • 4-3 量的是"结果对不对",这一节看"过程对不对"
  • 2-4-5 推出了两个端点的闭式解,它们正好可以当尺子
    • t = 0 t = 0 t=0 时 u 0 ( x ) = μ − x u_0(x) = \mu - x u0(x)=μ−x,其中 μ \mu μ 是数据均值
    • t = 1 t = 1 t=1 时 u 1 ( x ) = x u_1(x) = x u1(x)=x
  • 检查办法是在 − 3 , 3 2 -3, 3^2 −3,32 上取网格,逐点算 v θ ( x , t ) v_\theta(x, t) vθ(x,t),再与解析解相减
  • 在 t = 0 t = 0 t=0,模型在整个网格上都指向数据均值,中位误差不足 0.1。下图把五个时刻的场画在一起,灰蓝色的点是该时刻真实的边缘分布 p t p_t pt,绿色箭头是模型学到的 v θ v_\theta vθ
  • 从 t = 0 t = 0 t=0 到 t = 0.75 t = 0.75 t=0.75,场很平滑,箭头都指向数据所在的两条弯月,这就是模型在把概率质量往数据上搬
  • t = 0 t = 0 t=0 那一格的场几乎是纯放射状。那时还没有任何数据信息可用,唯一的线索就是均值,式 u 0 ( x ) = μ − x u_0(x) = \mu - x u0(x)=μ−x 说的正是这件事
  • 到 t = 0.25 t = 0.25 t=0.25 与 t = 0.5 t = 0.5 t=0.5,场里已经能看出两个吸引中心,两个月的分岔就是从这里开始的
  • t = 1 t = 1 t=1 那一格不一样,要单独说
4-4-1 t 接近 1 时速度场为什么发散
  • 先看数据附近。按 u 1 ( x ) = x u_1(x) = x u1(x)=x 验证,取样几个点,模型输出都贴着 x x x 本身
  • 但把整个网格铺开量, t = 1 t = 1 t=1 的中位误差比 t = 0 t = 0 t=0 大了两个数量级,近两成的网格点速度幅值超过 20,尾部一直拖到 150 以上
  • 原因在条件速度的闭式里: u t ( x ∣ x 1 ) = ( x 1 − x ) / ( 1 − t ) u_t(x|x_1) = (x_1 - x)/(1 - t) ut(x∣x1)=(x1−x)/(1−t),分母在 t → 1 t \to 1 t→1 时趋于零
    • 训练时 t t t 是均匀采样的,抽到恰好等于 1 的概率为零,所以训练数据里根本没有这一格
    • 网络在 t t t 接近 1 且远离数据的地方没有任何约束,只能外推。那里没有数据告诉它该往哪走,输出就发散
  • 逐时刻看速度幅值的中位数更直观
t t t 0 0.25 0.5 0.75 1.0
∣ v ∣ \lvert v \rvert ∣v∣ 中位数 2.51 2.66 3.11 4.57 3.27
  • 中间几个时刻的场是有界的,量级在 2 到 5 之间。 t = 1 t = 1 t=1 那一档的中位数看着不高,是因为大部分网格点仍落在正常范围,但它拖着一根长尾
  • 这件事没有影响本次采样,因为轨迹在 t t t 接近 1 时已经落在数据附近,而那里的场是正常的。但边界是清楚的:这里的发散来自 CondOT 这条路径、均匀时间采样与网格外推区三者的叠加,不是 Flow Matching 的普遍性质,所以时间采样分布不能随便往端点推
  • 轨迹图如下。左边是 48 条从噪声到数据的轨迹,颜色表示时刻;右边是这些轨迹的终点落在哪
  • 轨迹从四周的噪声区出发,先朝中间收拢,再分岔进两条弯月。个别轨迹走的是绕路,但终点都落在数据上
  • 右边的落点图给出的结论与 4-3 一致:终点铺满了两条弯月,拐角与两端的收尖处也没有空缺
4-5 改变 ODE 积分步数 ------ 少步采样换来多少质量、省下多少时间
  • 前三节都在看"能不能生成"。这一节换一个轴:采样要跑几次网络
  • 把 N N N 从 1 扫到 128,每一步记两个量
    • 生成质量:能量距离,在 4096 个样本上取多种子均值
    • 单样本耗时:4096 个样本并成一个批量跑一次完整采样,重复多次取最小值,再除以样本数摊薄到单个样本
  • 两个统计量的取法都是被实测逼出来的
    • 用多个种子而不是一个:同一档步数换种子,能量距离能差到 2 倍,单种子曲线会出现假拐点
    • 用最小值而不是中位数:单次读数只有几毫秒,启动开销与同步抖动同量级。任何外来干扰只会让某一次变慢而不会变快,最小值才是"不被干扰时到底多快"的估计
  • 完整读数如下
N N N(网络求值次数) 能量距离(多种子均值) 标准误 单样本耗时(毫秒) 相对 N = 128 N = 128 N=128 加速
1 0.562 1.9 × 10 − 4 1.9 \times 10^{-4} 1.9×10−4 0.000147 100.8 倍
2 0.163 4.2 × 10 − 4 4.2 \times 10^{-4} 4.2×10−4 0.000268 55.2 倍
4 0.0369 4.1 × 10 − 4 4.1 \times 10^{-4} 4.1×10−4 0.000498 29.7 倍
8 0.00664 2.35 × 10 − 4 2.35 \times 10^{-4} 2.35×10−4 0.000968 15.3 倍
16 0.00201 1.71 × 10 − 4 1.71 \times 10^{-4} 1.71×10−4 0.001887 7.85 倍
32 0.00110 1.74 × 10 − 4 1.74 \times 10^{-4} 1.74×10−4 0.003887 3.81 倍
64 0.000894 1.71 × 10 − 4 1.71 \times 10^{-4} 1.71×10−4 0.007818 1.89 倍
128 0.000852 1.86 × 10 − 4 1.86 \times 10^{-4} 1.86×10−4 0.014801 1.00 倍
  • 表里的耗时摊薄到了单个样本,所以绝对值很小。换算回整批就便于核对: N = 128 N = 128 N=128 那一档是 4096 个样本、128 次网络求值共约 61 毫秒,摊到每次求值约 0.5 毫秒
  • 生成样本的样子按步数排开如下,每个面板里灰色是真实数据、绿色是模型生成
  • 质量与耗时随步数的变化画在一起,两条曲线都放在对数轴上
  • 三条读数
4-5-1 单步采样:代价最低,误差最大
  • N = 1 N = 1 N=1 时 Euler 只走一步, x 1 = x 0 + v θ ( x 0 , 0 ) x_1 = x_0 + v_\theta(x_0, 0) x1=x0+vθ(x0,0)。而 t = 0 t = 0 t=0 处的解析场正好是 u 0 ( x ) = μ − x u_0(x) = \mu - x u0(x)=μ−x
  • 代入就得到

x 1 = x 0 + μ − x 0 = μ x_1 = x_0 + \mu - x_0 = \mu x1=x0+μ−x0=μ

  • 所有样本都落在数据均值上 ,与 x 0 x_0 x0 无关。噪声被完全消掉了
  • 实测完全吻合:输出的均值与数据几乎相同,标准差却小了一个多数量级。图 30 的第一个面板就是这个紧团,能量距离也说明它离真实分布很远
  • 这一步不是模型没学好,而是一步 Euler 在这条路径上必然丢掉几乎所有信息 。噪声 x 0 x_0 x0 在公式里被消掉,多样性也就没了
  • 加到两步就恢复了:均值仍在数据均值附近,但标准差回到与数据同量级,图 30 第二个面板里已经能看出两个月的雏形

说人话:一步采样把"从哪个噪声出发"这件事完全抹掉了,于是所有样本挤在同一个点上

4-5-2 质量在哪里饱和
  • 从 N = 1 N = 1 N=1 到 N = 8 N = 8 N=8,每翻一倍步数,能量距离就掉一大截,这正是 3-6 说的全局误差主导区
  • 从 N = 16 N = 16 N=16 起,质量进入 10 − 3 10^{-3} 10−3 量级,之后收益迅速衰减
    • 16 到 32 改善约四成
    • 32 到 64 改善约两成
    • 64 到 128 改善不到一成
  • 平台的位置要从逐种子的差来判,不能只看均值差
    • 32 步与 64 步的均值差已经落在种子噪声之内,分不出来
    • 16 步与 32 步的均值差明显超出噪声,可以分辨
    • 所以"再翻倍还有没有用"这件事,在 64 步这里断了
  • 离地板还有多远: N = 128 N = 128 N=128 的读数与同分布地板仍差一个多数量级
    • 这不是步数不够造成的。步数从 64 翻到 128 只买回不到一成,说明剩下的差距来自别的地方
    • 来源是网络容量与有限样本估计的噪声,也就是 3-6 说的那个地板。4-2 的预算扫描指向同一件事:把训练步数翻十倍,距离也只好了约一半
  • 换个说法:在这个任务上,16 步之后每多跑一次网络,买到的改善都在快速贬值。图 31 左边那条曲线在对数轴上的斜率变化,量的就是这件事
4-5-3 代价是线性的,收益不是
  • 单样本耗时从 N = 1 N = 1 N=1 到 N = 128 N = 128 N=128,两端正好差 100 倍
  • 中间每一档的相邻比值都在 2 附近,步数翻倍,耗时翻倍,没有折扣
    • 两端的偏差来自固定的启动开销: N N N 很小时它占比大, N N N 很大时它占比小
    • 这条直线就是图 31 右边那张图,标题里的"代价对步数线性"说的正是它
  • 一边线性、一边快速衰减,两者放在一起就得到本节的主结论:本次实验里,性价比拐点在 8 到 16 步之间,不在 128 步。这个区间是二维双月形、这个容量的网络与 Euler 求解器共同给出的,换成更复杂的分布或更大的模型都要重新标定
  • 把几个可用工作点列出来,对照如下
N N N 耗时降到 质量差多少倍
1 1/100 约 660 倍
8 1/15 约 8 倍
16 1/8 约 2.4 倍
32 1/4 约 1.3 倍
≥ 64 \ge 64 ≥64 按步数线性上涨 进平台
  • N = 1 N = 1 N=1 那一档唯一的用途是验证流程通了; N = 8 N = 8 N=8 的轮廓已经成形,但离饱和还差一个量级; N = 16 N = 16 N=16 用七分之一的代价换来两倍多的距离差,是这一列里最划算的位置
  • 最后要说清楚这一节没有 证明的东西。它只扫了 CondOT 这一条路径,没有做路径之间的对照,所以这些读数不能用来宣称某条路径"更好"
    • 2-6 说过两条路径的差别在形状,而形状好不好走得用同一把尺子量才有意义
    • 更重要的是,这个任务太简单,几条路径在几十步之后会撞到同一个地板,撞上地板之后谁快谁慢就分不出来了
    • 少步数那一端(1 到 16 步)才是差异可能出现的地方,而要把那一端做成有意义的对照,需要比二维双月形更难的分布

总结

  • 本篇从 Diffusion Policy 的推理开销出发,把 Flow Matching 的推导、实现与实验走了一遍:第 2 章给出从连续性方程到条件流匹配损失的推导,第 3 章拆解官方实现的四个部件,第 4 章在二维双月形上量了质量与代价
  • 四条值得记住的读数
    • 训练目标可以简化到一行平方误差,因为论文证明了它与那个不可算的 Flow Matching 目标梯度等价
    • 损失曲线早就平了,生成质量却还在变好。判收敛要看生成质量本身
    • 速度场在两个端点有闭式解: t = 0 t = 0 t=0 时是 μ − x \mu - x μ−x, t = 1 t = 1 t=1 时是 x x x。前者在整个网格上都对得上,后者只在数据附近成立,远处会发散
    • 采样步数上,质量在 64 步进平台,代价全程线性,本次实验中性价比拐点在 8 到 16 步
  • 最想留下的一条是:Flow Matching 的价值不在于"生成得更好",而在于把每一步的代价与路径的形状解耦了。路径直,大跨步也能走对;路径弯,就只能用小碎步
  • 这也解释了它为什么对机器人动作生成特别有吸引力。控制回路只给几十毫秒,能不能压进这个窗口,全看质量饱和得早不早
  • 下一期我们会把这条路径接到机器人动作生成上,看它如何用几次网络求值预测一段动作序列,以及条件输入在这套架构里怎么接
  • 如有错误,欢迎指出!
  • 感谢观看!
相关推荐
把我格子衫拿来1 小时前
PnP 位姿估计:从二维像素求三维姿态,为什么 tvec 不是相机位置?
python·opencv·计算机视觉·pnp·相机位姿
benchmark_cc1 小时前
实时行情系统如何避免无意义的重复请求?从请求去重到缓存设计
大数据·开发语言·python·pandas·量化交易·股票数据·quantdash
田里的水稻1 小时前
Python_两大法宝dir() 与 help() 的妙用
开发语言·python
郝学胜-神的一滴1 小时前
Python浏览器自动化:Selenium与Playwright实战对比,网页爬虫&自动操作一站式指南
开发语言·爬虫·python·selenium·程序人生·自动化
Zip-List1 小时前
工作小记-pytorch的算子优化: 转gpu核函数 python调c
c语言·c++·pytorch·python·cuda
隐擎fox1 小时前
解密 BGP Anycast 边缘路由与跨国链路质量检测:多租户网络出口的低时延与防丢包架构实战
网络·python·网络协议·tcp/ip·网络安全·架构·tcpdump
把我格子衫拿来1 小时前
手眼标定:相机与机械臂如何建立坐标关系?从 AX=XB 到 OpenCV 实践
python·opencv·机器人·手眼标定·坐标变换
把我格子衫拿来2 小时前
相机标定误差很小,为什么测量仍不准?从 OpenCV 重投影误差到独立验证
python·数码相机·opencv·计算机视觉·相机标定
goehou2 小时前
流式输出实战:SSE 解析、chunk 拼装与断流重试
python·ai·llm·sse·流式