Computing All Optimal Partial p-Wasserstein Matchings on the Line

会议: NeurIPS 2026
arXiv: 2608.18875
代码: https://github.com/JacobusTheSecond/wasserstein
领域: 学习理论
关键词: 最优传输, p-Wasserstein距离, 部分匹配, 快速傅里叶变换, 计算几何

一句话总结

针对一维线段上部分 \(p\)-Wasserstein 最优传输（OT-profile）中 \(p>1\) 导致的二次复杂度瓶颈，本文基于线段单调配对结构提出了 FFT 加速的平衡区间匹配数据结构，首次在 \(O(p\,n\log^2 n)\) 近线性时间内计算出全部最优部分匹配代价，并给出了 \(p=\infty\) 依赖 \((\min,+)\)-卷积猜想的二次下界。

研究背景与动机

\(p\)-Wasserstein 距离是度量概率分布之间差异的基础几何工具，在计算机视觉、生成模型、分布鲁棒优化与切片 Wasserstein 距离中发挥着核心作用。在全传输（Full Optimal Transport）设定下，一维分布由于保持保序保序（Monotone Matching）特性，仅需将红蓝点集排序后按序一一配对，便可在 \(O(n\log n)\) 时间内完成精确求解。这一极端高效的计算特性不仅支撑了高维切片最优传输算法，也是一维统计分析的基石。

然而在实际数据场景中，观测数据普遍存在噪声、离群点或仅有部分重叠，直接强行全量匹配会显著扭曲真实的几何距离。为了寻找最优的质量保留比例与离群点剔除平衡，计算各截断规模 \(k\in[n]\) 下最小传输代价的全局最优传输曲线（OT-profile）成为鲁棒最优传输的核心需求。对于 \(p=1\)，Chapel 与 Tavenard 近期在 ICLR 2025 中利用 \(L_1\) 度量的度量结构和退化嵌套性质，实现了 \(O(n\log n)\) 的全局计算。但对于严格凸的 \(p>1\) 代价函数，这一退化性彻底消失：凸性使得最优匹配的边集必须全局保序，从 \(k\) 规模增广到 \(k+1\) 时，既有匹配边全部发生偏移重排，所有 \(n\) 种规模的最优匹配累积包含 \(\Theta(n^2)\) 条互不相同的匹配边，直接显式维护匹配会导致不可逾越的 \(\Omega(n^2)\) 计算瓶颈。

面对 \(p>1\) 是否存在二次计算壁垒这一理论悬念，本文敏锐地洞察到：尽管全局匹配边集在逐步移位，但其在几何拓扑上可紧凑折叠为至多 \(n\) 个互不相交的“平衡区间”，且增广路在拓扑上仅影响极局部的区间合并。核心 idea：将增广路径的代价评估转化为平衡区间的最优保序匹配查询，设计基于线段树与快速多项式卷积（FFT）的区间匹配数据结构，在 \(O(p\,n\log^2 n)\) 预处理后实现 \(O(\log n)\) 单次平衡增广查询，彻底攻克 \(p>1\) 下部分匹配的二次时间瓶颈。

方法详解

整体框架

本文算法继承匈牙利算法（Hungarian algorithm）的原始-对偶交替增广框架，利用一维直线上局部最优匹配无交叉、无跨越自由点的结构特性，维护当前匹配的“平衡区间紧凑表示”。在每一轮增广中，算法从空匹配开始，依次寻找净代价（Net Cost）最小的增广路径，将 \(k\)-匹配扩展为 \((k+1)\)-匹配。

整个算法的计算效率取决于能否快速找到净代价最小的增广路径。本文证明了最优增广路径的两个端点必然是一对“相邻自由点”（Adjacent Free Pair），且增广路径的净代价完全由包含该自由点对的外层区间与原内部平衡区间的差值决定。借助预先构建的平衡区间匹配数据结构，该查询在 \(O(\log n)\) 时间内即可完成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入点集 R 与 B<br/>排序并合并全局坐标序列"] --> B["预处理区间匹配数据结构<br/>线段树各节点按二项式展开并行执行 p+1 次 FFT 卷积"]
    B --> C["初始化相邻自由点对优先队列<br/>初始增广净代价设为几何间距距离"]
    C --> D["贪心选取最小净代价增广路径<br/>弹出堆顶相邻自由点对并翻转增广路径"]
    D --> E["更新平衡区间紧凑表示<br/>替换内部旧区间并合并相邻无自由点区间"]
    E --> F["动态更新相邻关系与区间代价差<br/>调用数据结构在 O(log n) 时间内评估新自由点对净代价"]
    F -->|规模 k 递增至 n| G["输出全部最优部分匹配代价序列<br/>获得完整的全局 OT-profile"]

关键设计

1. 增广路径的平衡区间紧凑表示与代价归约

在 \(p>1\) 的直线上，若直接显式维护匹配边的集合，随着规模 \(k\) 自增，所有配对均因保持严格保序而整体移动，累积边数达到 \(\Theta(n^2)\)。本文证明了一维凸代价下的几何隔离性质：在最小代价匹配中，匹配边的内部绝不会包含任何未匹配的自由点（Free Point）。若存在自由点落在匹配边内部，直接将匹配端点替换为该自由点将严格缩短边长，破坏最优性。

基于此，本文将由全部匹配边构成的连通分量进行极大合并，定义了匹配的平衡区间表示 \(\mathcal{I}_M\)。每个平衡区间内红蓝点数完全相等，匹配在区间内部完全退化为单调升序对齐（第 \(i\) 个蓝点匹配第 \(i\) 个红点）。对于任意相邻自由点对 \((b, r)\)，其所构成的区间 \(I_{b, r} = [\min\{b, r\}, \max\{b, r\}]\) 内部至多包含 \(\mathcal{I}_M\) 中的一个平衡区间 \(I\)。增广路径 \(P\) 的净代价 \(\Phi(b, r)\) 即可严格解析为纯粹的区间匹配代价之差： $$ \Phi(b, r) = \begin{cases} |b - r|^p, & \text{若 } I_{b, r} \text{ 内部不包含 } \mathcal{I}M \text{ 中的任何区间}, \[6pt] w^p(I}) - w^p(I), & \text{若 } I \in \mathcal{IM \text{ 包含在 } I $$ 这一关键性质将复杂的全局图论增广路径权值更新，完全归约为两次独立的平衡区间匹配代价查询 } \text{ 内部}. \end{cases\(w^p(\cdot)\)。

2. 线段树分解与偏移量解耦的区间代价查询

为了在次线性时间内计算任意平衡区间 \(I\) 的最优匹配代价 \(w^p(I)\)，本文构建了一棵覆盖全部 \(2n\) 个红蓝点有序合并索引的标准二分线段树 \(\mathcal{T}\)。由于平衡区间内部红蓝点严格单调匹配，设 \(I\) 中最左侧红点全局编号为 \(j\)、最左侧蓝点全局编号为 \(k\)，则整个区间内部所有配对的编号偏移量（Shift）均为常数 \(d = k - j\)。

对于线段树中的任意节点 \(S\)，其管辖点集涵盖的可能偏移量范围为有限区间。当查询平衡区间 \(I\) 时，线段树将其递归分解为 \(O(\log n)\) 个标准区间： - 若节点 \(S \subseteq I\)，直接取用预存的节点内部匹配代价 \(W(S, d)\)； - 若区间 \(I\) 横跨节点 \(S\) 的左右子树 \(\mathrm{left}_S\) 与 \(\mathrm{right}_S\)，本文证明所有横跨左右子树的偏移量为 \(d\) 的点对必然全部完整落在 \(I\) 内部，绝无外部泄漏点。因此可精确拆解为： $$ \mathrm{Cost}(I \cap S, d) = W(\mathrm{left}_S, \mathrm{right}_S, d) + W(\mathrm{right}_S, \mathrm{left}_S, d) + \mathrm{Cost}(I \cap \mathrm{left}_S, d) + \mathrm{Cost}(I \cap \mathrm{right}_S, d) $$ 通过该递归合并公式，在 \(O(\log n)\) 树高内仅需常数次查表，即可严格求得 \(w^p(I)\)。

3. 二项式展开与多项式快速卷积加速预处理

为在不产生二次开销的前提下填入线段树各节点的跨子树配对总代价 \(W(\mathrm{left}_S, \mathrm{right}_S, d)\)，算法利用了线段上红蓝点的几何空间隔离。设左子树红点为 \(r_u, \dots, r_v\)，右子树蓝点为 \(b_s, \dots, b_t\)，由于左子树全部点在几何坐标上均严格位于右子树点左侧，距离绝对值可脱去外层绝对值符号，即 \(|b_{j+d} - r_j|^p = (b_{j+d} - r_j)^p\)。

利用二项式定理，将该代数式展开： $$ (b_{j+d} - r_j)^p = \sum_{c=0}^p \binom{p}{c} (-1)^c r_j^c \, b_{j+d}^{p-c} $$ 进而将跨区间总代价表示为 \(p+1\) 个离散卷积之和： $$ W(\mathrm{left}S, \mathrm{right}S, d) = \sum (-1)^c \left( \sum}^p \binom{p}{c \right) $$ 定义多项式 } r_j^c \, b_k^{p-c\(P_c(x) = \sum_{j=u}^v r_j^c x^{v - j}\) 与 \(Q_c(x) = \sum_{k=s}^t b_k^{p-c} x^{k - s}\)，其多项式乘积 \(P_c(x)Q_c(x)\) 在 \(x^\ell\) 处的卷积系数恰好对应于偏移量 \(d = \ell + s - v\) 的所有配对和。对于区间长度为 \(|S|\) 的节点，单次 FFT 仅需 \(O(|S|\log |S|)\) 时间，处理全部 \(p+1\) 个阶数仅耗时 \(O(p|S|\log |S|)\)。自底向上汇总后，全树构建时间被严格控制在 \(O(p\,n\log^2 n)\)，彻底规避了暴力枚举配对的 \(O(n^2)\) 瓶颈。

一个完整示例

考虑 \(n=2\)，\(p=2\) 的简易点集：红点位于坐标 \(R = \{1, 4\}\)，蓝点位于坐标 \(B = \{2, 7\}\)。 全局有序排列为 \(r_1=1 < b_1=2 < r_2=4 < b_2=7\)。 1. \(k=1\) 阶段：初始无匹配区间，优先队列中相邻自由对为 \((b_1, r_1)\) 距离差 \(|2-1|^2=1\)，\((b_1, r_2)\) 距离差 \(|2-4|^2=4\)，\((b_2, r_2)\) 距离差 \(|7-4|^2=9\)。堆顶弹出 \((b_1, r_1)\)，建立匹配边 \((b_1, r_1)\)，初始匹配代价 \(C_1 = 1\)。此时生成唯一平衡区间 \(I_1 = [1, 2]\)，剩余自由点为 \(R_F=\{r_2=4\}, B_F=\{b_2=7\}\)。 2. \(k=2\) 阶段：此时新的相邻自由对为 \((b_2, r_2)\)，对应区间 \(I_{b_2, r_2} = [4, 7]\)。由于其内部不包含已存在的区间 \(I_1=[1, 2]\)，其净代价计算直接采用单边代价 \(|7-4|^2=9\)。 若自由对跨越了 \(I_1\)，算法则调用数据结构分别查询并集大区间 \(w^2(I_{b_2, r_1})\) 与旧区间 \(w^2(I_1)\) 的差值。增广后得到完整匹配 \(C_2 = 1 + 9 = 10\)。数据结构仅需调用常数次 \(O(\log n)\) 树节点卷积系数求和，无需重排边集。

实验关键数据

主实验

论文开源了高效 C++ 实现 FFTHungarian，与 Chapel & Tavenard (ICLR 2025) 的高精度 Numba/njit 加速基线 PAWL 以及朴素区间匹配基线 naiveHungarian 在不同数据分布和样本规模下进行运行时间评测（测试平台：Apple M4，24GB RAM，设置 60s 运行超时阈值）。下表展示在 \(p=2\) 下不同分布差异度与数据量级的运行时间对比（单位：秒）。

分布族与配置	数据规模 \(n\)	FFTHungarian (本文)	PAWL (ICLR'25)	naiveHungarian	相对加速比
同分布高斯 \(\sigma_1=\sigma_2=1.0\)	10,000	0.042	0.048	0.31	1.14×
同分布高斯 \(\sigma_1=\sigma_2=1.0\)	100,000	0.58	0.62	> 60.0 (超时)	1.07×
异方差高斯 \(\sigma_2=1.1\) (微小差异)	10,000	0.045	0.18	1.85	4.00×
异方差高斯 \(\sigma_2=1.1\) (微小差异)	100,000	0.61	16.4	> 60.0 (超时)	26.9×
异方差高斯 \(\sigma_2=2.0\) (显著差异)	10,000	0.048	0.42	4.20	8.75×
异方差高斯 \(\sigma_2=2.0\) (显著差异)	50,000	0.29	10.5	> 60.0 (超时)	36.2×
异方差高斯 \(\sigma_2=2.0\) (显著差异)	100,000	0.65	> 60.0 (超时)	> 60.0 (超时)	> 92.3×
异均值高斯 \(\mu=0.1\)	100,000	0.62	22.1	> 60.0 (超时)	35.6×
单峰对双峰 \(\mu=0.5\)	100,000	0.64	38.5	> 60.0 (超时)	60.2×

消融实验

为了在工程实践中兼顾小规模区间的常数开销与大规模区间的渐近效率，本文引入了混合阈值截断策略：对包含少于 \(M\) 个点的平衡区间采用朴素快速求和，对大于 \(M\) 个点的区间调用线段树查询。下表展示在固定 \(n=100,000\)、\(\sigma_2=1.1\) 条件下不同截断阈值 \(M\) 的运行耗时与内存开销分析：

截断阈值配置 \(M\)	整体执行耗时 (s)	数据结构构建占比	动态增广查询占比	内存占用峰值 (MB)	设计说明
\(M=0\) (纯 FFT 树查询)	0.94	62%	38%	84.2	极小区间线段树递归常数开销偏大
\(M=512\)	0.73	64%	36%	84.2	小区间切换为向量化直接累加
\(M=2048\) (本文选用)	0.61	68%	32%	84.2	查询与递归开销最佳平衡点
\(M=8192\)	0.89	68%	32%	84.2	大区间朴素累加渐近退化开始显现
\(M=\infty\) (完全不建树，纯朴素)	16.4	0%	100%	12.5	完全退化为 PAWL 式二次复杂度瓶颈

关键发现

• 分布偏离引发二次退化：当两分布完全相同时，增广路径极短且高度局部化，基线 PAWL 表现良好；但只要分布均值或方差存在哪怕 10% 的微小偏离，跨区间长程增广迅速增多，基线算法立即退化至 \(\Theta(n^2)\) 复杂度，耗时急剧攀升。
• 近线性扩展优势显著：FFTHungarian 的耗时曲线在双对数坐标轴上始终保持斜率约为 1 的近线性增长，在 \(n=100,000\) 且分布存在差异的实际场景下，取得数十倍乃至接近两个数量级的加速，且是唯一能在 1 分钟时限内处理十万点级异质分布的方法。

亮点与洞察

• 将几何单调性转化为代数卷积：巧妙利用一维最优传输保序性导致的固定偏移量（shift）特征，将区间距离多项式通过二项式展开，化解为有限次离散多项式卷积，使得计算几何问题转化为经典信号处理算法。
• 精密的条件计算复杂度下界分离：论文不仅给出了上界算法，还针对 \(p=\infty\) 建立了归约证明，指出在 \((\min,+)\)-卷积猜想下不可能存在真次二次算法。这在理论上严谨界定了全传输（线性对齐即可求解）与部分传输之间的内在计算难度鸿沟。
• 软硬件兼顾的工程混合架构：算法并未盲目崇拜渐近复杂度，而是结合硬件缓存与小规模向量化特性设置 \(M=2048\) 的混合阈值，使理论上的近线性算法在工业级小常数实现中落地。

局限与展望

• 复杂度对阶数 \(p\) 的线性依赖：预处理开销包含因子 \(p\)，对于高阶矩（例如 \(p \ge 10\)）或非整数实数 \(p\)，二项式展开不再有限截断，需采用泰勒展开或多项式逼近，存在精度损失与阶数膨胀。
• 仅适用于一维几何空间：核心的保序性与区间覆盖无交叉性质严格依赖一维直线的全序拓扑结构，无法直接推广到二维或更高维度的几何空间。
• 未来方向：探索高维切片部分最优传输（Sliced Partial Optimal Transport, SPOT）中快速调用本文算法的端到端集成，以及针对带权分布（非均匀概率质量）的高效近似求解方案。

相关工作与启发

• vs PAWL (Chapel & Tavenard, ICLR 2025): PAWL 仅在 \(p=1\) 下利用 \(L_1\) 范数退化性实现了 \(O(n\log n)\)，而在 \(p>1\) 严格凸时代价评估陷入 \(\Theta(n^2)\) 循环；本文基于 FFT 构造的平衡区间数据结构，首次将任意有限整数 \(p\ge 1\) 的计算复杂度全部压缩至 \(O(p\,n\log^2 n)\)。
• vs 匈牙利算法与高维最近点对 (Agarwal et al., 1999): 一般几何度量下的增广轨线算法依赖动态双色最近点对维护，即便在低维也需要 \(\widetilde{O}(n^2)\) 耗时；本文彻底利用一维线段上的拓扑有序性，实现了与全传输几乎同量级的近线性效率。

评分

• 新颖性: ⭐⭐⭐⭐⭐ [攻克了一维部分最优传输长年遗留的二次瓶颈，FFT 结构设计极具启发性]
• 实验充分度: ⭐⭐⭐⭐☆ [理论与合成数据验证扎实，若能补充真实图像切片传输下游任务将更为完备]
• 写作质量: ⭐⭐⭐⭐⭐ [数学推导严密自洽，从定理证明到附录下界归约结构极度清晰]
• 价值: ⭐⭐⭐⭐⭐ [为切片鲁棒最优传输、部分概率匹配提供了关键的基础理论与高效实用工具]

相关论文

• [ICLR 2026] Slicing Wasserstein over Wasserstein via Functional Optimal Transport
• [NeurIPS 2026] Geometry-Aware Optimal Transport: Fast Intrinsic Dimension and Wasserstein Distance Estimation
• [ICML 2025] Theoretical Performance Guarantees for Partial Domain Adaptation via Partial Optimal Transport
• [NeurIPS 2026] Hyper Input Convex Neural Networks for Shape Constrained Learning and Optimal Transport
• [NeurIPS 2026] On Making SE(2)-Invariant Networks Optimal