损失指导流匹配

前置

• 流匹配

损失指导的流匹配

损失指导流匹配（Loss-Guidance Flow Matching，LGFM）在无条件向量场上加入损失梯度，使生成轨迹倾向于减小指定目标函数。

设无条件场为 $u_t$，损失为
$l_t(\mathbf{x},{\mathbf{x}}_1)$，受指导场定义为

$$v_t(\mathbf{x})=u_t(\mathbf{x})-{\nabla }_{\mathbf{x}}{l}_t(\mathbf{x},{\mathbf{x}}_1).\text{\hspace{1em}(1)}$$

这里 ${\mathbf{x}}_1$ 是指导条件，在以下关于 $\mathbf{x}$ 的微分中视为固定参数。

沿流的密度变化

连续性方程 (Villani, 2008) 为

$${\partial }_t{p}_t(\mathbf{x})+\mathrm{div}(p_t(\mathbf{x})v_t(\mathbf{x}))=0.\text{\hspace{1em}(2)}$$

若 $p_t>0$，且
${\varphi }_t^v$ 是 $v_t$ 的轨迹，

$$\frac{\mathrm{d}}{\mathrm{d}t}{\varphi }_t^v({\mathbf{x}}_0)=v_t({\varphi }_t^v({\mathbf{x}}_0)),{\varphi }_0^v({\mathbf{x}}_0)={\mathbf{x}}_0,$$

则链式法则把式 (2) 化为瞬时换元公式 (Ben-Hamu et al., 2022)

$$\frac{\mathrm{d}}{\mathrm{d}t}\log p_t^v({\varphi }_t^v({\mathbf{x}}_0))+\mathrm{div}{v}_t({\varphi }_t^v({\mathbf{x}}_0))=0.\text{\hspace{1em}(3)}$$

为保证以下等式逐点成立，可假设 $u_t$ 关于空间至少 $C^1$，
$l_t$ 关于空间至少 $C^2$，二者满足局部 Lipschitz 与保证
$[0,1]$ 上全局解存在的增长条件，并且 $p_t^u,p_t^v$ 是正的
$C^1$ 密度。正则性较弱时，相应结论应在弱解或几乎处处意义下解释。

由

$$\mathrm{div}{v}_t=\mathrm{div}{u}_t-\Delta l_t,$$

式 (3) 给出精确的轨迹积分公式

$$\begin{array}{rl}\log p_1^v({\varphi }_1^v({\mathbf{x}}_0))-\log p_0({\mathbf{x}}_0)=& -{\int }_0^1\mathrm{div}{u}_t({\varphi }_t^v({\mathbf{x}}_0))\mathrm{d}t\\ & +{\int }_0^1\Delta l_t({\varphi }_t^v({\mathbf{x}}_0),{\mathbf{x}}_1)\mathrm{d}t.\end{array}\text{\hspace{1em}(4)}$$

等价地，

$$\begin{array}{rl}{p}_1^v({\varphi }_1^v({\mathbf{x}}_0))=p_0({\mathbf{x}}_0)& \exp \left\{-{\int }_0^1\mathrm{div}{u}_t({\varphi }_t^v({\mathbf{x}}_0))\mathrm{d}t\right\}\\ \times & \exp \left\{{\int }_0^1\Delta l_t({\varphi }_t^v({\mathbf{x}}_0),{\mathbf{x}}_1)\mathrm{d}t\right\}.\end{array}\text{\hspace{1em}(5)}$$

式 (5) 保留了原推导中“散度—拉普拉斯—曲率”的关系，但第一指数因子一般不能直接认作无条件终点密度
$p_1^u({\varphi }_1^v({\mathbf{x}}_0))$：无条件密度的换元公式应沿
$u_t$ 自己的轨迹 ${\varphi }_t^u$ 积分，而这里积分路径是受指导轨迹
${\varphi }_t^v$。

与无条件密度之比的严格公式

令 $p_t^u$ 与 $p_t^v$ 分别满足

$$\begin{array}{rl}{\partial }_t{p}_t^u+\mathrm{div}(p_t^u{u}_t)& =0,\\ {\partial }_t{p}_t^v+\mathrm{div}(p_t^v{v}_t)& =0,\end{array}{p}_0^u=p_0^v=p_0.\text{\hspace{1em}(6)}$$

沿受指导轨迹计算无条件密度，有

$$\begin{array}{rl}\frac{\mathrm{d}}{\mathrm{d}t}\log p_t^u({\varphi }_t^v)& =\left({\partial }_t+v_t\cdot \nabla \right)\log p_t^u\\ & =-\mathrm{div}{u}_t-\nabla l_t\cdot \nabla \log p_t^u.\end{array}\text{\hspace{1em}(7)}$$

式 (3) 则给出

$$\frac{\mathrm{d}}{\mathrm{d}t}\log p_t^v({\varphi }_t^v)=-\mathrm{div}{u}_t+\Delta l_t.\text{\hspace{1em}(8)}$$

两式相减，得到密度比的输运 PDE

$$\left({\partial }_t+v_t\cdot \nabla \right)\log \frac{p_t^v}{p_t^u}=\Delta l_t+\nabla l_t\cdot \nabla \log p_t^u.\text{\hspace{1em}(9)}$$

因此精确的终点关系是

$$\boxed{\log \frac{p_1^v({\varphi }_1^v({\mathbf{x}}_0))}{p_1^u({\varphi }_1^v({\mathbf{x}}_0))}={\int }_0^1\left[\Delta l_t+\nabla l_t\cdot \nabla \log p_t^u\right]({\varphi }_t^v({\mathbf{x}}_0))\mathrm{d}t}.\text{\hspace{1em}(10)}$$

这里初始密度比为 $1$，所以没有初值项。式 (10) 表明，若要把受指导终点密度写成“无条件密度乘拉普拉斯权重”，还会出现指导梯度与无条件 score 的配对项。

原先常见的写法

$$p_1^v({\varphi }_1^v)\approx p_1^u({\varphi }_1^v)\exp \left\{{\int }_0^1\Delta l_t({\varphi }_t^v)\mathrm{d}t\right\}\text{\hspace{1em}(11)}$$

只能作为下列情形中的启发式近似：

• 精确条件：
  $\nabla l_t\cdot \nabla \log p_t^u=0$
  沿几乎所有受指导轨迹、几乎处处成立；
• 近似条件：遗漏项的轨迹积分足够小，即 $$\left|{\int }_0^1\nabla l_t\cdot \nabla \log p_t^u({\varphi }_t^v)\mathrm{d}t\right|\ll 1.$$

更一般地，若该绝对值不超过 $\epsilon$，则式 (11) 在对数密度比上的误差不超过 $\epsilon$，乘性误差落在
$[e^{-\epsilon },e^{\epsilon }]$ 内。没有这些 PDE 或近似条件时，沿受指导轨迹直接套用无条件密度的推导一般不成立。

梯度、拉普拉斯与曲率

对 $f:{\mathbb{R}}^d\to \mathbb{R}$，梯度为

$$\nabla f={\left(\frac{\partial f}{\partial x_1},\frac{\partial f}{\partial x_2},\dots ,\frac{\partial f}{\partial x_d}\right)}^{\top }.$$

拉普拉斯是梯度的散度：

$$\Delta f=\nabla \cdot (\nabla f)=\sum _{i=1}^d\frac{{\partial }^{2}f}{\partial x_i^2}=\mathrm{tr}{\nabla }^{2}f.\text{\hspace{1em}(12)}$$

${\partial }_{x_i}^{2}f$ 描述 $x_i$ 方向的局部曲率，
$\Delta f$ 是 Hessian 各特征方向曲率之和。在二阶可微的局部极小点，
${\nabla }^{2}f⪰0$，故 $\Delta f\ge 0$；在局部极大点，
${\nabla }^{2}f⪯0$，故 $\Delta f\le 0$。只有 Hessian 正定或负定时，才分别
得到严格的 $\Delta f>0$ 或 $\Delta f<0$。严格极值本身不保证严格符号，
例如 $f(x)=-x^4$ 在 $0$ 处是严格局部极大点但 $f^{\prime \prime }(0)=0$。
反过来，$\Delta f>0$ 只约束曲率之和，不保证每个方向都是“谷底”。

由式 (5)，正的累计拉普拉斯会增大沿同一受指导流的 Jacobian 密度因子，负值会减小它。若比较受指导密度和无条件密度，则必须使用式 (10)，同时计入
$\nabla l_t\cdot \nabla \log p_t^u$；因此不能只凭拉普拉斯符号断言最终样本密度一定被放大或抑制。

Ben-Hamu, H., Cohen, S., Bose, J., Amos, B., Grover, A., Nickel, M., Chen, R. T. Q., & Lipman, Y. (2022). Matching Normalizing Flows and Probability Paths on Manifolds. arXiv Preprint arXiv:2207.04711.

Villani, C. (2008). Optimal Transport: Old and New (Vol. 338). Springer.