Computing All Optimal Partial \(p\)-Wasserstein Matchings on the Line¶
会议: NeurIPS 2026
arXiv: 2608.18875
OpenReview: https://openreview.net/forum?id=88SUwJD17z
代码: https://github.com/JacobusTheSecond/wasserstein
领域: 其他方向
关键词: 最优传输, 偏Wasserstein距离, 快速傅里叶变换, 匈牙利算法, 条件下界
一句话总结¶
针对一维实数线上偏 \(p\)-Wasserstein 匹配轮廓计算长期面临的 \(\Omega(n^2)\) 复杂度瓶颈，本文提出一种基于二项式展开与快速傅里叶变换（FFT）卷积的平衡区间线段树数据结构，将全体最优偏匹配序列的求解时间突破至 \(O(p n \log^2 n)\)，并证明了 \(p=\infty\) 情形下基于 \((\min, +)\)-卷积猜想的二次时间下界。
研究背景与动机¶
Wasserstein 距离作为衡量概率分布间几何差异的核心工具，在机器学习生成建模、形状分析与统计推断中得到了极其广泛的应用。在全量最优传输设定下，一维分布间的最优匹配展现出极强的几何简捷性：只需将红蓝两组点集沿实数轴排序后按序对齐，即可在线性或近线性时间 \(O(n \log n)\) 内求出全局最优解，这一优雅特性也是切片 Wasserstein（Sliced Wasserstein）等大规模近似算法的基石。然而，在实际数据分析中，输入往往伴随着严重噪声、离群点或部分重叠，强行传输全部概率质量会造成对异常值的高度敏感。为此，偏最优传输（Partial Optimal Transport）与非平衡最优传输备受关注，其核心在于求解最优传输轮廓（OT-profile）——即针对所有传输质量比例 \(\alpha \in (0, 1]\)（离散设定下即所有基数 \(k \in \{1, \dots, n\}\) 的 \(k\)-匹配），分别求出最小传输代价 \(C_k\)。
对于 \(p=1\) 的特殊情形，Chapel 与 Tavenard（ICLR 2025）曾证明利用距离度量的退化嵌套结构，从最优 \(k\)-匹配转移到 \((k+1)\)-匹配只需新增单条匹配边而不破坏已有边，从而在 \(O(n \log n)\) 时间内算出全谱轮廓。但对于任何 \(p > 1\)，目标函数的严格凸性彻底打破了这种嵌套性：即便最优 \(k\)-匹配选取的仍然是排序后的前缀与后缀，为了最小化凸代价，所有已匹配点都必须重新按序对齐。这就导致当匹配规模从 \(k\) 增至 \(k+1\) 时，整组匹配边会发生全局整体平移，前后两步的匹配边集合可能没有任何交集。在 \(n\) 步增广过程中，可能产生 \(\Theta(n^2)\) 条完全不同的匹配边。以往显式维护匹配边或基于传统增广轨增量更新的算法（如经典匈牙利算法特化版本 PAWL）均不可避免地跌入 \(\Omega(n^2)\) 的时间泥潭，导致大规模分布的高阶偏传输匹配在计算上不可行。
面对高阶偏传输在一维空间中看似必然出现的二次时间壁垒，本文的核心切入点是：虽然不同基数下的显式匹配边集发生了剧烈变动，但一维最优匹配的几何本质必然保持局部的保序单调性，其全局状态完全可以被极少量的连通“平衡区间”隐式捕捉。核心 idea：将增广路径的代价更新规约到平衡区间上的最优对齐代价差值，并构建基于多项式二项式展开与快速傅里叶变换（FFT）的线段树数据结构，仅需 \(O(p n \log^2 n)\) 预处理与 \(O(n \log n)\) 空间即可在 \(O(\log n)\) 时间内响应任意平衡区间的最优匹配代价查询，从而彻底突破高阶偏 Wasserstein 轮廓的二次计算瓶颈。
方法详解¶
整体框架¶
算法的整体流水线建立在一维排序点集的原始-对偶匈牙利增广框架之上。给定实数轴上的红点集合 \(R\) 与蓝点集合 \(B\)（大小均为 \(n\)），算法首先将全部 \(2n\) 个点合并排序，并在排好序的点序列上构建一棵平衡区间匹配代价线段树 \(\mathcal{D}\)。整个计算过程从空匹配 \(M_0 = \emptyset\) 开始，历经 \(n\) 次增广迭代，依次构造出规模为 \(k=1, \dots, n\) 的最优偏匹配 \(M_k\)。
在每次迭代中，算法不需要显式追踪每一条具体的点对匹配边，而是动态维护匹配点集的外包平衡区间集合 \(\mathcal{I}_M\) 以及自由点（未匹配点）构成的相邻空闲对优先队列 \(\mathcal{Q}\)。算法从堆中取出净代价最小的相邻红蓝自由点对，执行单步增广；增广操作仅引起常数个平衡区间的插入与相邻合并，并产生至多一个新增的相邻自由点对。随后，算法通过向线段树 \(\mathcal{D}\) 发起至多两次 \(O(\log n)\) 时间的区间匹配查询来精确更新该点对的增广代价。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入点集 R, B 排序合并<br/>构建长度 2n 的全局顺序"] --> B["二项式展开与多项式卷积<br/>构建线段树 D (时间 O(pn log² n))"]
B --> C["初始化相邻自由点对优先队列 Q<br/>维护稀疏平衡区间集合 I_M"]
C --> D["选出当前最小净代价增广对 (b, r)<br/>从堆 Q 弹出最优相邻自由点对"]
D --> E["动态更新平衡区间与空闲点集<br/>合并相邻区间并移除已匹配点"]
E --> F["对数时间平衡区间代价查询<br/>查询 D 计算新邻接对净代价 Φ(b', r')"]
F -->|迭代 k = 1 到 n| G["输出全部最优 k-偏 Wasserstein 匹配代价<br/>生成完整 OT-Profile C_1, ..., C_n"]关键设计¶
1. 紧凑区间表示与增广路径代数刻画：将全局匹配状态坍缩为极少区间
为了解决匹配边随 \(k\) 增大而全盘更替导致的 \(\Theta(n^2)\) 边膨胀问题，本文利用了一维最优匹配的关键几何性质：任何最优匹配边构成的闭区间内部绝不可能包含未匹配的自由点（否则缩短该边必然严格降低总代价）。据此，算法将匹配边跨越的连通分支提炼为不相交的“平衡区间”集合 \(\mathcal{I}_M\)（每个区间内红点与蓝点数量严格相等），并在相邻区间之间无自由点时实施贪心合并。这一表示具有极强的数学性质：任意相邻自由点对 \((b, r)\) 构成的闭区间 \(I_{b,r}\) 内部至多包含 \(\mathcal{I}_M\) 中的一个平衡区间。因此，以 \((b, r)\) 为端点的最小增广路径 \(P\) 的净代价 \(\Phi(P)\) 可以精炼地表示为区间代价的差分：若 \(I_{b,r}\) 内无已有区间，则净代价直接等于 \(|b-r|^p\)；若包含已有区间 \(I \in \mathcal{I}_M\)，则净代价严格等于扩充后大区间的全局保序匹配代价与内部旧区间代价之差：
$$
\Phi(b, r) = w^p(I_{b,r}) - w^p(I)
$$
这就将复杂的全局增广轨搜索完全规约为了平衡区间上的单调匹配代价计算。
2. 基于二项式展开与 FFT 卷积的平衡区间线段树：打破 \(O(n^2)\) 预处理瓶颈
若每次计算 \(w^p(I)\) 均需遍历区间点集，单次增广仍需耗费线性时间。针对任意平衡区间 \(I\)，其内部红点与蓝点必按单调顺序严格匹配，若 \(I\) 中首个红点全局索引为 \(j\)、首个蓝点全局索引为 \(k\)，则所有匹配边具有恒定的索引偏移（shift）\(d = k - j\)。为了快速响应任意平衡区间在任意合法偏移下的代价，作者构建了一棵覆盖 \([1, 2n]\) 的平衡线段树 \(\mathcal{T}\)。在线段树的每个节点 \(S = [a, b]\)，红点落在左子树 \(\mathsf{left}_S\)、蓝点落在右子树 \(\mathsf{right}_S\) 的跨越代价满足 \(b_{j+d} > r_j\)。通过二项式定理将幂次项彻底展开：
$$
|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} $$ 固定 \(c\) 后，跨越项代价在所有相关偏移 \(d\) 上的求和式 \(\sum_{j} r_j^c b_{j+d}^{p-c}\) 正好对应于序列 \(A = (r_u^c, \dots, r_v^c)\) 与翻转序列 \(B = (b_t^{p-c}, \dots, b_s^{p-c})\) 的离散卷积 \((A * B)\)。利用快速傅里叶变换（FFT），单个节点可在 \(O(p |S| \log |S|)\) 时间内并行计算出所有可能偏移跨越代价的系数多项式。自底向上归纳后，整棵树的构建时间被严格控制在 \(O(p n \log^2 n)\)，占用空间仅为 \(O(n \log n)\)。
3. 对数时间平衡区间代价查询与区间合并更新：实现 \(O(\log n)\) 步长快速增广
当查询一个平衡区间 \(I\) 的最优匹配代价 \(w^p(I)\) 时，算法首先通过二分查找确定 \(I\) 中最左侧红点与蓝点，直接算出其唯一偏移值 \(d\)。随后在线段树中进行标准的单点/区间覆盖递归分解：若当前节点 \(S \subseteq I\)，直接返回预存的子树全匹配代价 \(W(S, d)\)；若 \(I\) 同时跨越 \(S\) 的左右子树，则由于区间内部匹配的保序性与平衡性，本文严密证明了所有在左右子树之间跨越的 shift-\(d\) 边必定完全落在 \(I\) 内部（不会漏算也不会多算外部点对），因此只需累加预存的跨越代价 \(W(\mathsf{left}_S, \mathsf{right}_S, d) + W(\mathsf{right}_S, \mathsf{left}_S, d)\) 并递归访问左右两子树。由于线段树区间分解仅产生至多两条根至叶的查询路径，查询过程访问的节点数为 \(O(\log n)\)，单次平衡区间代价查询仅需 \(O(\log n)\) 时间。每次增广仅发生常数次堆调整与至多 2 次区间查询，使得全部 \(n\) 步增广累计仅需 \(O(n \log n)\) 时间，全局复杂度牢牢由预处理时间 \(O(p n \log^2 n)\) 支配。
损失函数 / 训练策略¶
作为离散组合优化算法，本方法没有深度神经网络的连续梯度反向传播过程，其优化目标是极小化各阶离散偏匹配代价：
$$
C_k = \min_{M \subset B \times R, \, |M|=k} \sum_{(b,r) \in M} |b - r|^p, \quad \forall k \in {1, \dots, n}
$$
在工程实现与工程优化策略上，作者给出了以下关键设计：
- 小规模区间朴素计算阈值截断：在线段树查询开销与朴素单调遍历之间存在常数平衡点。作者在开源 C++ 实现 FFTHungarian 中设置了大小为 2048 的阈值：当待查询平衡区间包含的点数少于 2048 时，直接采用 \(O(M)\) 的线性保序累加；仅当区间点数大于等于 2048 时才调用线段树的 FFT 预计算系数，兼顾了大区间的理论渐近优势与小子区间的极低硬件常数。
- 条件计算复杂性边界划分：论文针对 \(p = \infty\) 给出了严格的不可行性证明。通过将任意序列上的 \((\min, +)\)-卷积规约到单调序列上的 \((\max, -)\)-卷积，进一步规约至 \(k\)-偏 \(\infty\)-Wasserstein 匹配轮廓计算，证明了在 \((\min, +)\)-卷积猜想下，不存在能在 \(O(n^{2-\varepsilon} \mathrm{polylog} W)\) 时间内求解 \(p=\infty\) 偏匹配轮廓的真次二次算法。这表明算法依赖于有限 \(p\) 的二项式展开是本质性的，揭示了一维全量最优传输（\(O(n \log n)\) 对任意 \(p \ge 1\) 均成立）与偏最优传输之间的深刻计算复杂性断层。
实验关键数据¶
主实验¶
实验在配备 Apple M4 处理器和 24 GB 内存的平台进行，对比了本文提出的 C++ 实现 FFTHungarian、Chapel & Tavenard（ICLR 2025）经 Numba/njit 高度优化的 SOTA 算法 PAWL 以及纯 Python 朴素增广基线 naiveHungarian（\(p=2\)，单次运行超时上限设定为 60 秒）。合成数据覆盖不同方差高斯对 \(\mathcal{N}(0, 1^2)\) 与 \(\mathcal{N}(0, \sigma^2)\)、不同均值高斯对 \(\mathcal{N}(0, 1^2)\) 与 \(\mathcal{N}(\mu, 1^2)\)、以及单峰与双峰混合高斯。下表汇总了中等差异合成场景下的中位数耗时表现：
| 实验场景分布设置 | 样本规模 \(n\) | naiveHungarian | PAWL (ICLR 2025 SOTA) | FFTHungarian (本文) | 加速比 (vs PAWL) |
|---|---|---|---|---|---|
| 高斯均值偏移 \(\mu = 0.1\) | 1,000 | 0.42 s | 0.03 s | 0.04 s | 0.75× (常数开销相当) |
| 高斯均值偏移 \(\mu = 0.1\) | 5,000 | 12.80 s | 0.85 s | 0.22 s | 3.86× |
| 高斯均值偏移 \(\mu = 0.1\) | 10,000 | > 60 s (超时) | 4.12 s | 0.49 s | 8.41× |
| 高斯均值偏移 \(\mu = 0.1\) | 25,000 | > 60 s | 31.50 s | 1.35 s | 23.33× |
| 高斯均值偏移 \(\mu = 0.1\) | 50,000 | > 60 s | > 60 s (超时) | 2.91 s | > 20× (破局二次瓶颈) |
| 高斯标准差偏移 \(\sigma = 1.1\) | 10,000 | > 60 s | 3.94 s | 0.46 s | 8.57× |
| 高斯标准差偏移 \(\sigma = 1.1\) | 50,000 | > 60 s | > 60 s (超时) | 2.80 s | > 21× |
| 双峰混合偏移 \(\mu = 0.5\) | 10,000 | > 60 s | 4.60 s | 0.52 s | 8.85× |
| 双峰混合偏移 \(\mu = 0.5\) | 50,000 | > 60 s | > 60 s (超时) | 3.10 s | > 19× |
消融实验¶
为了剖析算法各模块的有效性与理论边界，作者开展了理论下界分离、数据结构查询类型以及工程混合阈值划分的系统性分析：
| 算法机制 / 问题设定 | 预处理时间复杂度 | 单次查询 / 增广复杂度 | 全谱 \(n\) 步总时间复杂度 | 条件限制 / 理论瓶颈 |
|---|---|---|---|---|
| 单平衡区间查询 (本文 FFT 线段树) | \(O(p n \log^2 n)\) | \(O(\log n)\) | \(O(p n \log^2 n)\) | 要求有限整数 \(p \in \mathbb{N}\) |
| \(p = \infty\) 极大值偏匹配 (本文下界) | \(\Omega(n^2 / 2^{\sqrt{\Omega(\log n)}})\) | — | \(\Omega(n^{2-\varepsilon})\) | \((\min, +)\)-卷积猜想 |
| 双任意区间匹配查询 (两区间问题) | 组合算法下界 | 组合算法下界 | \(P \cdot Q \in \Omega((n_B n_R)^{1-\varepsilon})\) | 组合矩阵乘法猜想 |
| 双任意区间匹配查询 (非组合算法) | 代数矩阵乘法 | 代数矩阵乘法 | \(P + n Q \in \Omega(n^{\omega/2})\) | 矩阵乘法指数 \(\omega\) 限制 |
| FFTHungarian (全朴素区间计算) | \(O(1)\) | \(O(M)\) 遍历 | \(\Theta(n^2)\) | 仅分布完全重合时退化较轻 |
| FFTHungarian (混合阈值 \(M < 2048\)) | \(O(p n \log^2 n)\) | \(O(\min(M, \log n))\) | \(O(p n \log^2 n)\) | 小区间免除 FFT 解卷积常数 |
关键发现¶
- 对非同分布输入的决定性优势：当且仅当红蓝两分布完全相同时（\(\sigma=1.0, \mu=0\)），增广路径极少改变历史匹配边，PAWL 尚能维持近线性；然而现实应用中分布总是存在差异（甚至仅偏移 10%），PAWL 迅速退化到最坏二次耗时 \(\Theta(n^2)\)，在 \(n \ge 50,000\) 时均在 60 秒内超时，而本文算法始终保持平滑的近线性扩展。
- 理论边界的深刻分离：一维单平衡区间查询（本问题核心）允许 \(O(p n \log^2 n)\) 近线性求解，但一旦放宽到不相交的“任意双区间匹配”，便撞上组合矩阵乘法下界（需 \(\Omega(n^2)\)）；而将度量提升至 \(p=\infty\) 时，又撞上 \((\min, +)\)-卷积的真二次下界。这充分印证了本文设计的精巧性正好处于能够突破二次屏障的黄金理论窗口中。
亮点与洞察¶
- 多项式代数与线段树的完美结合：传统计算几何难以处理非度量性质的 \(p > 1\) 凸成本；本文巧妙利用二项式定理把 \(|b - r|^p\) 拆成幂次乘积求和，将几何对齐转化为 FFT 离散卷积，极富启发性。
- 状态稀疏化表征消除二次边爆炸：明确指出一维偏匹配即便伴随 \(\Theta(n^2)\) 条相异边，其拓扑结构仍能压缩为至多 \(n\) 个不相交平衡区间，为高维传输向一维切片投影后的偏匹配算法提供了强有力的工具。
- 理论闭环与完备性：不仅给出了有限整数 \(p\) 的近线性算法，还建立起了 \(p=\infty\) 的条件强下界，完整界定了该问题在计算复杂性图谱上的坐标。
局限与展望¶
- 复杂度与幂次 \(p\) 呈线性正比：虽然对常用的 \(p=1, 2, 3\) 计算极快，但预处理时间正比于 \(p\)；对于高阶非整数矩 \(p \notin \mathbb{N}\)，无法直接采用有限项二项式展开，仍需借助多项式逼近或级数截断。
- 点权重必须为均匀质量 \(1/n\)：算法严重依赖于点集离散均匀假设下的有序一对一匹配；对于任意实数质量权重的非均匀离散点集，保序索引偏移将演变为连续质量切片，算法扩展存在非平凡的技术挑战。
相关工作与启发¶
- vs Chapel & Tavenard (ICLR 2025, PAWL)：PAWL 仅在 \(p=1\) 时达到近线性，在 \(p>1\) 时退化为二次时间；本文在线段树与 FFT 支撑下首次在 \(p>1\) 领域将时间压缩到近线性。
- vs 经典匈牙利算法与高维最近点对算法：一般图增广需要 \(O(n^3)\)，低维空间需要 \(O(n^2)\)；本文针对一维几何特征实现了 \(O(p n \log^2 n)\)，可作为切片偏最优传输（Sliced Partial OT）的快速底层核算工具。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次打破 \(p > 1\) 一维偏 Wasserstein 轮廓的二次计算瓶颈，理论贡献卓越。
- 实验充分度: ⭐⭐⭐⭐ 综合对比了合成高斯分布族及双峰混合分布，开源 C++ 代码扎实，若能补充真实切片 Wasserstein 下游任务效果更佳。
- 写作质量: ⭐⭐⭐⭐⭐ 几何性质引理严密，数据结构设计逻辑清晰，上下界分析浑然一体。
- 价值: ⭐⭐⭐⭐⭐ 为离群点鲁棒学习、正无标签学习（PU Learning）及切片偏传输等算法提供了高效的底层算法支撑。