Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation¶
会议: NeurIPS 2026
arXiv: 2605.20122
OpenReview: https://openreview.net/forum?id=5GGreyp28W
领域: 优化/理论
关键词: Wasserstein距离, 计算-统计运行时, 抽样-草图-求解, 最优传输, 最小费用流
一句话总结¶
本文提出面向连续概率分布二次 Wasserstein 距离估计的「抽样-草图-求解」范式，在对样本建立笛卡尔网格中心投影草图后转化为多部图最小费用流精确求解，首次在 \(d=2\) 且 \(\alpha>1/2\) 的 Hölder 平滑分布下实现了匹配统计下界的 \(\tilde{O}(\varepsilon^{-2})\) 最优计算-统计运行时。
研究背景与动机¶
二次 Wasserstein 距离 \(W_2^2(P,Q)\) 是现代机器学习与统计分析中衡量连续概率分布差异的核心工具，广泛应用于最小距离估计（MDE）、非参数两样本检验以及生成模型评测。然而，在现有实际应用中，未知分布 \(P, Q\) 往往只能通过随机抽样获取经验样本点集，并调用离散最优传输算法求解经验测度间的距离。纯计算视角的离散最优传输算法（如匈牙利算法、Sinkhorn 迭代及其加速变体、图伸缩网络流算法等）通常将离散样本规模 \(n\) 和逼近精度视为独立参数，在二维或三维欧氏空间中随着样本量 \(n\) 的增加计算复杂度迅速膨胀，往往达到 \(O(n^2/\varepsilon)\) 甚至更高，且忽略了样本本身作为总体测度经验估计所固有的统计抽样误差。
现有方法在联合考量计算复杂度与统计误差时面临严峻的权衡困境：单纯增加经验样本量虽能压低统计误差，却使下游离散传输求解开销难以承受；而引入熵正则化（Entropic OT）或 Sinkhorn 发散虽能在计算上实现线性或拟线性迭代，但引入的正则化偏差要求极小的平滑参数，反噬了联合收敛速率，导致计算-统计运行时（Computational-Statistical Runtime, CSR）在低维空间中仍退化为 \(O(\varepsilon^{-6})\) 至 \(O(\varepsilon^{-7.5})\)，显著劣于直接经验估计的理论极限。其核心矛盾在于：缺乏一种既能利用总体分布平滑性大幅压缩样本规模、又不对传输几何结构引入不可控几何畸变的结构化离散化机制。
本文从计算-统计联合优化的全新视角切入，认为对连续分布距离估计的算法设计应始于对总体分布平滑性的先验刻画，并将样本收集成本（每次抽取开销 \(O(1)\)）直接计入总运行时间。核心 idea：提出「抽样-草图-求解」（Sample-Sketch-Solve）范式，基于有界网格中心点投影的对称一阶误差抵消机制对经验样本建立正规笛卡尔网格草图，将传输问题等价转化为有容量限制的多部图最小费用流精确求解，在不损失统计渐近精度的同时将总体计算复杂度逼近理论最优下界 \(\tilde{\Theta}(\varepsilon^{-2})\)。
方法详解¶
整体框架¶
本文构建的 Sample-Sketch-Solve（抽样-草图-求解）算法流程包含三个紧密耦合的阶段： 1. 抽样阶段（Sample）：从待估计的未知连续概率分布 \(P\) 与 \(Q\) 中分别独立采集 \(n = O(\varepsilon^{-\max(2, d/2)})\) 个随机样本，构建经验概率测度 \(\hat{P}_n\) 与 \(\hat{Q}_n\)； 2. 草图压缩阶段（Sketch）：将支撑域 \((0,1)^d\) 沿每个坐标轴等距划分为 \(L = 1/h = \Theta(\varepsilon^{-1/(1+\alpha)})\) 个区间，构建正规笛卡尔网格 \(\mathcal{G}_h\)，并将落入每个网格单元的所有样本质量坍缩投影到该网格的几何中心点 \(\bar{G}_h^d\)，生成稀疏离散直方图测度 \(G_h\#\hat{P}_n\) 与 \(G_h\#\hat{Q}_n\)； 3. 网络流求解阶段（Solve）：利用欧氏平方距离在正规网格上的坐标可分离性，将网格直方图之间的 \(W_2^2\) 计算严格归约为 \((d+1)\)-部图上的无容量最小费用流问题，通过整型伸缩后采用确定性拟线性最小费用流算法直接输出精确最优解 \(\hat{A}(\varepsilon)\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["连续分布 P, Q (0,1)^d<br/>alpha-Hölder 平滑且有界"] --> B["阶段 1：自适应统计抽样<br/>抽取 n=O(eps^-max(2,d/2)) 样本"]
B --> C["阶段 2：正规网格中心投影草图<br/>网格步长 h=O(eps^(1/(1+alpha)))<br/>中点对称性消除一阶离散误差"]
C --> D["阶段 3：多部图最小费用流精确求解<br/>坐标分离映射至 (d+1)-部图网络流<br/>整型缩放后调用拟线性算法求解"]
D --> E["精确输出 A_hat(eps)<br/>满足期望误差 E|A_hat - W_2^2| <= eps"]关键设计¶
1. 自适应统计抽样：匹配最优经验收敛界的样本量确定
传统离散最优传输将样本规模 \(n\) 视为给定输入，导致下游计算复杂度无法与统计估计目标自适应对齐。本文首先根据连续有界域上的经验经验测度收敛率理论，确定达到 \(\varepsilon/2\) 统计期望误差所需的样本规模。在 \(d < 4\) 维度下，经验测度逼近真实测度的均方误差为 \(O(n^{-1/2})\)，因而取 \(n = O(\varepsilon^{-\max(2, d/2)})\) 即可保证 \(\mathbb{E}|W_2^2(P,Q) - W_2^2(\hat{P}_n, \hat{Q}_n)| \leq \varepsilon/2\)。由于每次抽样开销为 \(O(1)\)，抽样阶段总计算时间严格控制在 \(O(\varepsilon^{-\max(2, d/2)})\) 范围内，为后续草图压缩提供了坚实且紧致的离散经验起点。
2. 正规网格中心投影草图：中点对称性诱导的一阶误差消除机制
直接将连续分布或密集点集量化离散到网格通常会导致 \(O(h)\) 的三角不等式上界误差，若要控制在 \(\varepsilon\) 误差内，需划分极细的网格（\(h = O(\varepsilon)\)），使得网格顶点数 \(L^d = O(\varepsilon^{-d})\) 极大，破坏计算可行性。本文的关键突破在于发现：当分布密度具有 \(\alpha\)-Hölder 平滑性时，将质量投影到网格几何中心能够利用类似数值积分中点法则（Midpoint Rule）的奇偶对称性抵消一阶误差项。具体而言，基于 Brenier 定理及 Collins-Tong 边界正则性结果，最优传输映射对应的位移场 \(\ell_\mu(x) = x - T_\mu(x)\) 满足 \((1-\xi)\)-Hölder 连续。在展开网格投影测度与真实测度的距离差时，交叉内积项被拆解为网格单元 \(C_j\) 上的积分：
由于 \(C_j\) 关于中心点 \(c_j\) 高度对称，线性项 \(x - c_j\) 的积分严格为零。剩余误差项完全由位移场差值和密度扰动的 Hölder 连续性控制，从而严格证明了离散化逼近误差界：
这一指数从 \(h\) 显著提升至 \(h^{1+\alpha}\)，使得只需选择较粗的网格步长 \(h = O(\varepsilon^{1/(1+\alpha)})\)（对应单轴划分数 \(L = \Theta(\varepsilon^{-1/(1+\alpha)})\)），即可保证离散误差不超过 \(\varepsilon/2\)，将草图离散点规模极致压缩至 \(L^d = O(\varepsilon^{-d/(1+\alpha)})\)。
3. 多部图最小费用流精确求解：坐标可分离性驱动的拟线性网络流转化
尽管草图投影后支撑集点数缩减至 \(L^d\)，但各网格单元的经验质量权重为非均匀离散值（均为 \(1/n\) 的整数倍），直接求解离散最优传输若使用传统稠密 LP 或 Sinkhorn 仍需 \(O((L^d)^2)\) 复杂度。本文利用欧氏平方距离在正规笛卡尔网格上的轴对称可分离性，即点坐标差异满足 \(\|x-y\|_2^2 = \frac{1}{L^2}\sum_{\ell=1}^d (a_\ell - b_\ell)^2\)，将高维稠密直方图匹配问题归约为一个具有 \((d+1)L^d\) 个节点、仅 \(m = d L^{d+1}\) 条有向弧的 \((d+1)\)-部图最小费用流问题。为彻底消除传统连续松弛的额外近似误差，本文对网络结构施加整数伸缩变换：
由于总供给量和需求量放大 \(n\) 倍后变为正整数，将弧容量设为 \(U \leq n\)，弧费用放大 \(L^2\) 倍后变为满足 \(C < L^2\) 的非负整数。随后直接调用 van den Brand 等人的确定性近线性时间最小费用流算法，在 \(\tilde{O}(m^{1+o(1)} \log U \log C) = \tilde{O}(L^{d+1+o(1)} \log n \log L)\) 时间内完成精确组合求解，彻底摆脱了近似算法对误差容限的依赖。
损失函数 / 训练策略¶
本算法属于非参数统计估计与离散组合优化流程，不涉及神经网络参数训练与梯度反向传播。核心求解目标为带整型供求平衡约束的有向最小费用流问题：
其中节点不平衡量 \(b_v = n \cdot (p_v - q_v) \in \mathbb{Z}\)，弧费用 \(c_e \in \{0, 1, \dots, (L-1)^2\} \subset \mathbb{Z}_{\geq 0}\)。算法全流程无需调参，网格分辨率 \(L\) 与样本量 \(n\) 均完全由预设精度目标 \(\varepsilon\) 和分布平滑先验阶数 \(\alpha\) 解析确定。
实验关键数据¶
主实验¶
论文对有界支撑域 \((0,1)^d\) 上满足 \(\alpha\)-Hölder 平滑与有界密度假设的概率分布，在达到 \(\varepsilon\)-加性期望估计误差下，对各类算法的计算-统计运行时（CSR，已忽略 \(\mathrm{poly}(\log(1/\varepsilon))\) 与 \(\varepsilon^{-o(1)}\) 因子）进行了系统理论对比。
| 方法 / 算法族 | 理论依据与特性 | CSR 通式 (\(d\) 维) | \(d=2\) 空间复杂度 | \(d=3\) 空间复杂度 |
|---|---|---|---|---|
| Sinkhorn 发散 (Chizat et al., 2020) | 熵正则化平滑对偶求解 | \(\varepsilon^{-(2\lfloor d/2 \rfloor + 5.5)}\) | \(\varepsilon^{-7.5}\) | \(\varepsilon^{-7.5}\) |
| Sinkhorn + Richardson 外推 (Chizat et al., 2020) | 多重正则化外推偏差消除 | \(\varepsilon^{-\frac{2\lfloor d/2 \rfloor + 11}{2}}\) | \(\varepsilon^{-6.5}\) | \(\varepsilon^{-6.5}\) |
| 熵正则化最优传输 (Chizat et al., 2020) | 原始 Entropic OT | \(\varepsilon^{-\max(6, d+2)}\) | \(\varepsilon^{-6.0}\) | \(\varepsilon^{-6.0}\) |
| 经验插件 + 离散图伸缩 (Lahn et al., 2019) | 经验采样 + 离散图逼近 | \(\varepsilon^{-(\max(d,4)+1-\mathbb{I}(d=2))}\) | \(\varepsilon^{-4.0}\) | \(\varepsilon^{-5.0}\) |
| 本文非平滑基线 (Corollary 4.1.3) | 正规网格草图 + 最小费用流 (\(\alpha=0\)) | \(\varepsilon^{-\max(2, d+1)}\) | \(\varepsilon^{-3.0}\) | \(\varepsilon^{-4.0}\) |
| 本文平滑范式 (Theorem 4.1) | \(\alpha\)-Hölder 草图 + 最小费用流 | \(\varepsilon^{-\max(2, \frac{d+1}{1+\alpha})}\) | \(\varepsilon^{-2.0}\) (\(\alpha > 1/2\)) | \(\varepsilon^{-2.0}\) (\(\alpha \to 1\)) |
消融实验¶
论文围绕分布平滑度指标 \(\alpha\)（对应于离散化误差界中 \(h^{1+\alpha}\) 的指数提升程度）以及网格中心对称投影对误差抵消的作用进行了理论消融分析：
| 理论配置 / 设定条件 | 离散化逼近误差界 | \(d=2\) 下的 CSR | \(d=3\) 下的 CSR | 理论机制与说明 |
|---|---|---|---|---|
| 完整框架 (Full: 中心投影 + \(\alpha\)-Hölder) | \(O(h^{1+\alpha})\) | \(O(\varepsilon^{-2})\) (当 \(\alpha > 0.5\)) | \(O(\varepsilon^{-2})\) (当 \(\alpha \to 1\)) | 中心对称消除一阶误差，达到理论最优统计率 |
| 退化平滑性设定 (\(\alpha = 0.5\)) | \(O(h^{1.5})\) | \(O(\varepsilon^{-2.0})\) | \(O(\varepsilon^{-2.67})\) | 在二维空间刚好达到最优统计下界 \(\varepsilon^{-2}\) |
| 弱平滑性设定 (\(\alpha = 0.2\)) | \(O(h^{1.2})\) | \(O(\varepsilon^{-2.5})\) | \(O(\varepsilon^{-3.33})\) | 连续加速，性能始终严格优于经典经验插件法 |
| 无平滑性保证 (\(\alpha = 0\) / 任意测度) | \(O(h)\) | \(O(\varepsilon^{-3.0})\) | \(O(\varepsilon^{-4.0})\) | 仅依赖几何有界性，无法触发一阶积分中点消项 |
| 非中心点任意投影 (w/o Center Projection) | \(O(h)\) | \(O(\varepsilon^{-3.0})\) | \(O(\varepsilon^{-4.0})\) | 破坏正规网格中心奇偶对称性，退化为传统三角不等式界 |
关键发现¶
- 统计最优复杂度的突破：在二维空间（\(d=2\)）中，只要分布平滑度 \(\alpha > 1/2\)（例如 Lipschitz 连续密度 \(\alpha=1\)），本文算法的整体计算-统计运行时严格达到 \(\tilde{\Theta}(\varepsilon^{-2})\)，与 Niles-Weed & Rigollet 证明的均绝对误差统计下界 \(\Omega(\varepsilon^{-2})\) 完全匹配，证明了在此类平滑测度上最优传输估计的时间开销可以完全被样本采集开销主导，计算几乎无需付出额外渐近代价。
- 三维空间的极限逼近：在三维空间（\(d=3\)）中，当密度函数逼近 Lipschitz 连续（\(\alpha \to 1\)）时，CSR 渐近收敛至 \(\varepsilon^{-\frac{4}{1+\alpha}} \to \varepsilon^{-2}\)，将传统最优估计速率从 \(\varepsilon^{-5}\) 巨幅压缩至理论极限边缘。
- 正规网格与中心投影的双重威力：单纯使用网格分割并不能带来阶数飞跃；只有将质量严格投影到网格中心点，使线性位移积分项在对称区域上恒为零，才能从传统 \(O(h)\) 三角不等式界跨越至 \(O(h^{1+\alpha})\)，这是支撑整个低采样规模草图的数学核心。
亮点与洞察¶
- 范式创新：从纯计算向计算-统计联合分析的范式转移：传统理论将抽样后的离散点集视为绝对精确输入，导致算法盲目追求离散点集的高精度传输；本文确立的 CSR 视角将样本收集、统计波动与下游计算融为一体，明确了计算精度无需超越统计误差的根本准则。
- 机理巧妙：中点数值积分对称性向最优传输误差分析的精妙迁移：巧妙利用 Brenier 映射和全局位移场的 Hölder 连续性，将网格投影误差分解为高阶残差与关于中心对称抵消的线性项，突破了长期依赖测度弱收敛或纯度量三角不等式的局限。
- 算法桥梁：连续流形几何与离散拟线性最小费用流的无缝衔接：通过坐标分解将高维直方图匹配精准转化为多部图网络流，并利用确定性拟线性网络流求解器实现无精度损失的组合求解，避免了连续凸优化或熵正则化带来的收敛精度损耗。
局限与展望¶
- 维度灾难的固有约束：当前基于正规笛卡尔网格的草图技术顶点数为 \(L^d\)，随着空间维度 \(d \ge 4\) 的增加，其计算复杂度的指数级依赖仍未消除，难以直接扩展到高维图像特征表示或高维深度嵌入空间。
- 紧支撑与边界平滑性假设：理论推导依赖于支撑集在超立方体 \((0,1)^d\) 内以及密度上下界严格远离零（\((\alpha,C)\)-controlled），在无界支撑或含有低维流形奇异性的实际数据分布上缺乏直接理论覆盖。
- 未来改进方向：可进一步探索自适应四叉树/八叉树层次网格或低维流形投影草图机制，在保持中心对称积分消除特性的同时，突破空间维数对网格划分规模的硬性限制。
相关工作与启发¶
- vs 熵正则化最优传输与 Sinkhorn 发散 (Chizat et al., 2020)：Chizat 等人开创了统计-计算权衡的理论分析，但由于熵正则化对 \(W_2^2\) 的逼近需要极小的正则参数，导致 Sinkhorn 迭代次数随精度迅速爆炸（在二维下 CSR 为 \(\varepsilon^{-6.5} \sim \varepsilon^{-7.5}\)）；本文摒弃连续对偶迭代，采用中心网格草图配合离散精确网络流求解，实现了真正的阶数跃迁（二维下达到 \(\varepsilon^{-2}\)）。
- vs 离散图伸缩逼近算法 (Lahn et al., 2019)：Lahn 等人针对固定离散样本点集给出了目前最好的几何近似算法（二维耗时 \(O(n/\varepsilon^2)\)），但若直接用于总体估计，其总耗时仍受制于 \(O(\varepsilon^{-4})\)；本文通过主动进行结构化网格投影与坐标解耦，大幅降低了图规模，比该离散基线整整快两个数量级。
- vs 经验测度逼近统计极限 (Weed & Bach, 2019; Niles-Weed & Rigollet, 2022)：统计文献证明了无偏经验测度估计的最小样本复杂度下界为 \(\Omega(\varepsilon^{-2})\)，但未给出多项式甚至近线性时间内实现该下界的实用算法；本文首次在理论上闭环证明了平滑分布在 \(d=2\) 下该统计下界在计算上是完全可达的。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出连续分布 Wasserstein 估计的 Sample-Sketch-Solve 范式并给出匹配统计下界的紧致复杂度分析。
- 实验充分度: ⭐⭐⭐⭐ 理论推导完备自洽，严格证明了离散化误差抵消与网络流映射复杂度，缺少真实经验合成数据集的数值实验（纯理论工作）。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑严谨清晰，从统计下界、中点对称误差抵消到多部网络流映射层层递进，定理证明精妙。
- 价值: ⭐⭐⭐⭐⭐ 解决了最优传输领域长期悬而未决的低维平滑分布计算-统计联合最优性难题，对统计两样本检验与生成评测具有重要理论指导意义。