势匹配最优传输：用于精确$p$-Wasserstein动力学的连续归一化流

# 用于精确 p-Wasserstein 动力学的连续归一化流

来源：https://arxiv.org/html/2608.05666

## 势匹配最优传输：用于精确 p-Wasserstein 动力学的连续归一化流

Lishuo Zhang¹, Ruizhi Huang¹, Yang Yu¹, Lei Li¹,²

1 上海交通大学数学科学学院  
2 上海交通大学自然科学研究院、教育部-科学计算重点实验室  
上海 200240，中国

[email protected]  
[email protected]  
[email protected]  
[email protected]  

###### 摘要

我们提出了势匹配最优传输（PMOT），一个面向一般 p-成本最优传输的势流框架，其中 \(c_p(x,y)=\|x-y\|^p\)。PMOT 以广义 Benamou–Brenier 形式（针对所选指数 \(p\)）用标量势参数化 CNF 速度场。它通过沿由模型自身端点确定的直线桥上的自诱导匹配损失来训练势梯度，同时允许灵活的终端分布匹配。我们的主要结果建立了零损失精确性：在所述正则性、精确终端匹配和唯一性假设下，任何零损失解都满足广义 Benamou–Brenier 最优性系统，并恢复相应的 p-最优传输映射和动力学。在合成基准上，PMOT 学习到的 p 特定映射与相应的 p-匹配 OT 参考一致。它作为基于似然的高维表格数据密度模型也保持竞争力，并且基于 MMD 的颜色变换实验展示了灵活的基于样本的终端匹配。

## 1 引言

连续归一化流（CNFs）通过时间相关的 ODE 动力学学习数据分布与易处理的参考先验之间的可逆变换（Chen et al., 2018；Grathwohl et al., 2018）。它们通过连续变量替换公式支持精确的似然估计，同时定义了分布之间的传输轨迹。这使得 CNF 成为基于似然的生成建模与最优传输（OT）之间的自然接口。

一条特别有影响力的研究路线通过将学习到的动力学正则化到低成本传输路径上，将 CNF 与 OT 联系起来。例如，OT-Flow（Onken et al., 2021）引入了势流参数化以及动能正则化，从而得到一种基于似然的模型，其动力学与二次成本最优传输的 Benamou–Brenier 公式密切相关（Benamou and Brenier, 2000）。然而，所得到的几何结构自然地与二次成本（即 \(W_2\) 几何）绑定。许多 OT 问题则是由更一般的成本 \(c_p(x,y)=\|x-y\|^p\) 定义的，其中指数 \(p\) 控制所诱导的传输几何，尤其是在非对称传输问题中，可能导致不同的最优耦合和轨迹（Villani and others, 2009; Peyré and Cuturi, 2019）。

这激励我们构建基于似然的势流，使其动力学不仅与二次情形对齐，也与一般 p-成本 OT 对齐。在这项工作中，我们引入了*势匹配最优传输*（PMOT），一个用于一般 p-成本最优传输并带有终端分布匹配的势流框架。当使用 KL/NLL 终端损失时，PMOT 保留了 OT-Flow 风格 CNF 的主要结构优势：动力学保持连续，并且似然通过 CNF 变量替换公式进行评估。PMOT 以共轭指数 \(q=p/(p-1)\) 在广义 Benamou–Brenier 形式中参数化速度场，从而使动力学依赖于目标 p-成本几何。与使用二次动能惩罚不同，PMOT 引入了自诱导的势匹配残差。对于每个数据样本 \(x\)，当前 CNF 产生一个终端端点 \(F_\theta(x)\)，残差将沿所诱导直线桥的势梯度与 \(-\|F_\theta(x)-x\|^{p-2}\bigl(F_\theta(x)-x\bigr)\) 匹配。该条件等价于将所诱导速度匹配到常速桥速度。它不需要预计算的 OT 耦合，而是对模型自身的端点映射施加一致性。

PMOT 的核心理论性质是零损失精确性结果。在所述正则性、精确终端匹配和唯一性假设下，每个零损失解都满足相应 p-成本问题的广义 Benamou–Brenier 最优性系统。因此，在这个理想极限中，所诱导的流恢复精确的 p-Wasserstein 动力学并达到最优传输成本。在有限样本的神经网络训练中，这种精确性受经验采样、模型容量、数值积分和优化误差的影响；因此，我们的经验性主张是用与 p-匹配 OT 参考的一致性来表述，而不是无条件的精确恢复。相应地，我们使用与训练模型相同指数计算的事后 OT 参考来评估几何保真度。

我们在合成传输、高维表格密度估计和图像颜色变换上评估 PMOT。实验评估了 p-匹配几何对齐、在温和目标加权下的生成、基于似然的密度建模以及基于 MMD 的样本传输。我们的贡献总结如下：

- 我们提出了 PMOT，一个用于一般 p-成本最优传输的势流 CNF 框架，通过自诱导势匹配进行训练，无需外部 OT 耦合或内部 OT 优化。
- 我们证明了，在所述正则性、精确终端匹配和唯一性假设下，每个零损失 PMOT 解都恢复相应的 p-最优传输映射和动力学。
- 我们使用以匹配成本指数计算的 OT 参考来评估每个模型，并展示了 p 特定的几何对齐、有竞争力的似然建模，以及通过 KL/NLL 和 MMD 目标实现的灵活终端匹配。

## 2 预备知识

### 2.1 一般 p-成本 OT 与位移插值

设 \(\mu_0,\mu_1\) 为 \(\mathbb{R}^d\) 上具有有限 p 阶矩的概率测度。对于 \(p\geq 1\)，p-成本最优传输问题为

\[
W_p^p(\mu_0,\mu_1)=\inf_{\pi\in\Pi(\mu_0,\mu_1)}\int\|x-y\|^p\,\mathrm{d}\pi(x,y).
\]

当最优 Monge 映射 \(T_p\) 存在时，它满足 \(T_{p\#}\mu_0=\mu_1\) 并达到相同成本。其位移插值为

\[
X_t(x)=(1-t)x+tT_p(x),\qquad \rho_t=(X_t)_{\#}\mu_0,
\]

速度为

\[
u_p(t,X_t(x))=T_p(x)-x.
\]

该路径达到动态 p-成本：

\[
\int_0^1\int_{\mathbb{R}^d}\|u_p(t,z)\|^p\,\mathrm{d}\rho_t(z)\,\mathrm{d}t
=
\int_{\mathbb{R}^d}\|T_p(x)-x\|^p\,\mathrm{d}\mu_0(x)
=
W_p^p(\mu_0,\mu_1).
\]

### 2.2 Benamou–Brenier 最优性与势速度

同一 p-成本传输问题具有动态 Benamou–Brenier 公式（Santambrogio, 2015）。动态问题为

\[
\inf_{\rho,v}\frac{1}{p}\int_0^1\int_{\mathbb{R}^d}\|v(t,x)\|^p\,d\rho_t(x)\,dt
\]

受端点约束 \(\rho_{t=0}=\mu_0\)，\(\rho_{t=1}=\mu_1\) 以及连续性方程

\[
\partial_t\rho_t+\nabla_x\cdot(\rho_t v)=0.
\]

在光滑的形式推导中，标量场 \(\Phi(t,x)\) 作为该连续性方程约束的 Lagrange 乘子出现。关于速度的相应 Euler–Lagrange 条件在 \(\rho_t\) 的支撑上给出

\[
\|v\|^{p-2}v=-\nabla_x\Phi.
\]

由于 \(p>1\)，映射 \(z\mapsto\|z\|^{p-2}z\) 是可逆的，因此最优速度具有势形式

\[
v=-\|\nabla_x\Phi\|^{q-2}\nabla_x\Phi=-\nabla_\xi H_q(\nabla_x\Phi),\qquad H_q(\xi)=\frac{1}{q}\|\xi\|^q.
\]

因此，PMOT 使用的势参数化并非任意假设：它反映了光滑 Benamou–Brenier 最优性系统中速度–乘子关系。关于 \(\rho\) 的 Euler–Lagrange 条件随后给出

\[
\partial_t\Phi+\nabla_x\Phi\cdot v+\frac{1}{p}\lVert v\rVert^p=0
\]

从而在 \(\rho_t\) 的支撑上有

\[
\partial_t\Phi-\frac{1}{q}\|\nabla\Phi\|^q=0.
\]

正如我们将看到的，这对应于粒子沿其轨迹速度保持恒定这一事实，这是 OT 的一个关键特征。在神经网络训练期间，\(\Phi_\theta\) 是一个模型参数；仅在理想零损失最优区域内，它才与 BB 乘子重合，且可相差一个仅依赖于时间的函数，这不会改变 \(\nabla_x\Phi_\theta\) 或所诱导的速度。

### 2.3 连续归一化流

连续归一化流通过常微分方程定义可逆映射

\[
\frac{dz_t}{dt}=v_\theta(t,z_t),\qquad z_0=x.
\]

设 \(F_\theta\) 表示时间一映射，\(F_\theta(x)=z_1\)。如果终端密度被选为易处理的先验，例如 \(\mathcal{N}(0,I)\)，则数据似然由瞬时变量替换公式计算：

\[
\frac{d}{dt}\log\rho_t(z_t)=-\nabla\cdot v_\theta(t,z_t).
\]

因此，

\[
\log\rho_0(x)=\log\rho_1(z_1)+\int_0^1\nabla\cdot v_\theta(t,z_t)\,dt,
\]

符号约定取决于积分方向。因此，CNF 既提供似然估计，又同时定义连续传输轨迹。

## 3 势匹配最优传输

我们引入*势匹配最优传输*（PMOT），一个用于一般 p-成本最优传输的变分框架。PMOT 通过标量势参数化 CNF 速度场，并优化两项：沿模型生成的直线桥的势匹配，以及终端分布匹配。在下面所述假设下，每个足够正则的零损失解都恢复唯一的 p-最优 Monge 映射及其相应的 Benamou–Brenier 速度场。

### 3.1 势诱导流

设 \(p>1\) 且 \(q=p/(p-1)\)。遵循上述 Benamou–Brenier 最优性关系，我们取标量势 \(\Phi:[0,1]\times\mathbb{R}^d\to\mathbb{R}\) 作为主变量，并定义所诱导的速度场

\[
v_\Phi(t,x)=-\lVert\nabla_x\Phi(t,x)\rVert^{q-2}\nabla_x\Phi(t,x). \tag{1}
\]

对于 \(p=2\)，方程 (1) 简化为熟悉的势流关系 \(v_\Phi(t,x)=-\nabla_x\Phi(t,x)\)。对于每个 \(x\in\mathbb{R}^d\)，考虑初值问题

\[
\partial_t z(t)=v_\Phi\bigl(t,z(t)\bigr),\text{ for almost every }t\in(0,1),\qquad z(0)=x. \tag{2}
\]

在整篇文章中，我们只限制在满足如下条件的势 \(\Phi\) 上：对每个 \(x\in\mathbb{R}^d\)，问题 (2) 在 \(\mathrm{AC}([0,1];\mathbb{R}^d)\) 中存在唯一解。我们将此解记为 \(z_\Phi(\cdot,x)\)。进一步要求相应的解映射 \(z_\Phi:[0,1]\times\mathbb{R}^d\to\mathbb{R}^d\) 是联合 Borel 可测的。

终端流映射 \(F_\Phi:\mathbb{R}^d\to\mathbb{R}^d\) 定义为

\[
F_\Phi(x)=z_\Phi(1,x).
\]

在时间 \(t\) 处诱导的分布为

\[
\rho_t^\Phi=\bigl(z_\Phi(t,\cdot)\bigr)_{\#}\mu_0.
\]

特别地，\(\rho_0^\Phi=\mu_0\)，\(\rho_1^\Phi=(F_\Phi)_{\#}\mu_0\)。

### 3.2 PMOT 变分目标

势匹配损失定义为

\[
\mathcal{L}_{\mathrm{PM}}(\Phi)
:=
\mathbb{E}_{x\sim\mu_0}\left[
\int_0^1
\left\|
\nabla_x\Phi\!\left(t,(1-t)x+tF_\Phi(x)\right)
+
\left\|F_\Phi(x)-x\right\|^{p-2}\left(F_\Phi(x)-x\right)
\right\|^2
\,\mathrm{d}t
\right].
\]

设 \(\mathcal{D}\) 为概率测度之间的非负差异，满足 \(\mathcal{D}(P,Q)=0\Leftrightarrow P=Q\)。我们定义终端分布损失为

\[
\mathcal{L}_{\mathrm{term}}(\Phi):=\mathcal{D}(\rho_1^\Phi,\mu_1).
\]

一个主要例子是前向 KL 散度，此时 \(\mathcal{L}_{\mathrm{term}}(\Phi)=\mathrm{KL}(\rho_1^\Phi\|\mu_1)\)。我们定义 PMOT 目标为

\[
\mathcal{L}_{\mathrm{PMOT}}(\Phi)
:=
\lambda_{\mathrm{PM}}\mathcal{L}_{\mathrm{PM}}(\Phi)
+
\lambda_{\mathrm{term}}\mathcal{L}_{\mathrm{term}}(\Phi),
\]

其中 \(\lambda_{\mathrm{PM}},\lambda_{\mathrm{term}}>0\) 是固定权重。定义域 \(\operatorname{dom}(\mathcal{L}_{\mathrm{PMOT}})\) 由所有满足前述流动适定性和联合 Borel 可测性要求且 \(\mathcal{L}_{\mathrm{PM}}(\Phi)\) 与 \(\mathcal{L}_{\mathrm{term}}(\Phi)\) 均有定义的势 \(\Phi\) 组成。

### 3.3 零损失蕴含最优传输

我们的主要定理是：

###### 定理 3.2（零损失下的精确恢复与速度场唯一性）

设 \(p\in(1,\infty)\)，设 \(\mu_0,\mu_1\in\mathcal{P}_p(\mathbb{R}^d)\)，并将 \(\operatorname{dom}(\mathcal{L}_{\mathrm{PMOT}})\) 记为 PMOT 目标的定义域。假设 \(\mu_0\) 关于 Lebesgue 测度绝对连续。假设 p-Benamou–Brenier 问题对某个 \(\Phi^\star\in\operatorname{dom}(\mathcal{L}_{\mathrm{PMOT}})\) 存在极小元 \((\rho^{\star},v^{\star})=(\rho^{\Phi^\star},v_{\Phi^\star})\)。则

\[
\min_{\Phi\in\operatorname{dom}(\mathcal{L}_{\mathrm{PMOT}})}\mathcal{L}_{\mathrm{PMOT}}(\Phi)
=
\mathcal{L}_{\mathrm{PMOT}}(\Phi^\star)
=
0.
\]

进一步地，设 \(\Phi\) 为一个全局极小元。如果 \(\Phi\in C^2((0,1)\times\mathbb{R}^d)\)，且如果 \(\operatorname{supp}(\rho_t^