A Unified Approach for Computing Wasserstein Barycenters of Discrete and Continuous Probability Measures¶

会议: NeurIPS 2026
arXiv: 2605.11270
OpenReview: https://openreview.net/forum?id=F23W2B2hf8
领域: 优化/理论
关键词: Wasserstein重心, 镜面对抗下降, 最优传输, Fisher-Rao几何, 半离散最优传输

一句话总结¶

提出在 Fisher-Rao 几何下利用 KL 散度作为 Bregman 距离的原始镜面下降算法 FRBary，以统一范式计算离散、连续及异构混合分布的精确无正则化 Wasserstein 重心，并在离散输入下自适应迭代生成具备理论收敛保证的绝对连续重心。

研究背景与动机¶

Wasserstein 重心作为将欧氏均值推广到测度空间的几何中心，通过最小化与一组输入概率测度之间的加权二次 Wasserstein 距离和，在分布式统计建模、分布对齐以及计算机图形学中具有核心地位。然而，其实用计算长久以来严重分化为离散与连续两套截然隔绝的计算体系：对于离散点云输入，经典方法依赖超大规模线性规划（LP），计算开销随支集点数立方级爆炸；对于连续测度，近年来基于对偶势函数的 Sobolev 梯度上升或神经网络参数化方法虽取得进展，却高度依赖规则网格差分或复杂的极小极大对抗训练，不仅无法直接处理无规则点云，也难以保证最优性条件的精确满足。

导致原始一阶梯度法难以取得通用理论收敛保障的根本症结在于：Wasserstein 重心目标泛函在测度空间原生的二次 Wasserstein 度量（$W_2$ 测地几何）下并不具备测地凸性。已有少数梯度流收敛保证多局限于高斯分布等特殊场景，或强行假设难以验证的 Wasserstein PŁ 不等式。更关键的是，许多下游统计任务（如分布回归、Wasserstein 主成分分析）迫切要求锚点重心本身具备平滑的绝对连续密度，但直接在离散点云上求解的离散重心本质上仍是离散测度，无法兼顾采样灵活性与连续密度建模。

本文的切入视角是跳脱出原生 $W_2$ 黎曼梯度的束缚，注意到重心目标泛函在测度的线性混合空间中具备严格凸性，其一阶变分可由各输入边际的最优传输 Kantorovich 势函数加权直接给出。核心 idea：将负相对熵作为镜像映射，在概率密度的 Fisher-Rao 黎曼几何下运行 Bregman 镜面下降，并在每步自适应调用半离散或离散最优传输求解器，构建能无缝统一离散、连续及混合输入的无正则化精确重心迭代算法。

方法详解¶

整体框架¶

算法以任意具备正下界界限的连续参考密度 $\rho^0$（如均匀分布）为起点，迭代过程中的当前重心估计始终保持为绝对连续测度 $\lambda^k \in \mathcal{P}_{\text{ac}}(\Omega)$。在每一轮迭代中，首先并行计算当前密度 $\rho^k$ 到各个输入测度 $\mu_i$ 的 Kantorovich 势函数 $\phi_{\rho^k \to \mu_i}$；随后依据输入测度类型（点云样本、网格直方图或连续密度），分别通过半离散传输的 Laguerre 胞腔划分或连续势函数求解完成一阶变分聚合；最后利用指数乘性权重更新当前连续重心密度，并重新做全局归一化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入分布集合<br/>点云/网格直方图/连续密度"] --> B["初始化参考密度<br/>紧支撑上的正下界密度"]
    B --> C["多类型自适应传输求解<br/>半离散Laguerre/网格离散OT"]
    C --> D["Kantorovich势加权变分<br/>一阶变分线性聚合"]
    D --> E["Fisher-Rao几何镜面更新<br/>对数域指数乘性迭代与归一化"]
    E -->|未达到收敛轮数 T| C
    E -->|收敛| F["输出连续Wasserstein重心<br/>显式连续密度或朗之万采样点云"]

关键设计¶

1. Fisher-Rao 几何下的连续密度镜面下降：规避测地非凸性的优化重构

针对传统 Wasserstein 梯度流在 $W_2$ 几何下非测地凸的理论死结，算法将优化空间转置到赋予 Fisher-Rao 黎曼度量的绝对连续概率测度空间 $\mathcal{P}_{\text{ac}}(\Omega)$。由于目标泛函 $\mathcal{E}(\lambda) = \sum_{i=1}^n \frac{w_i}{2} W_2^2(\mu_i, \lambda)$ 对密度混合具备线性凸性，其在 $\lambda$ 处的首次变分可精确表达为输入测度最优传输对偶势的加权和 $\delta \mathcal{E}_\lambda = \sum_{i=1}^n w_i \phi_{\lambda \to \mu_i}$。以负相对熵为镜像映射所导出的 Bregman 散度恰为 KL 散度 $D_{\text{KL}}(\rho \parallel \rho^k)$，从而导出了显式的指数乘性更新闭式解： $$\rho^{k+1}(x) = \frac{1}{Z_k} \rho^k(x) \exp\left( -\eta_k \sum_{i=1}^n w_i \phi_{\rho^k \to \mu_i}(x) \right)$$ 其中 $Z_k$ 为全域配分归一化常数。这一更新在形式上与常用的加熵正则化 Sinkhorn 重心有着本质区别：后者改变了原始的最优传输距离定义导致估计有偏，而本文仅将熵作为优化几何的镜面映射，优化的目标函数依然是无偏、精确的原始 Wasserstein 距离。

2. 半离散 Laguerre 胞腔自适应势函数求解：跨模态异构输入的统一兼容

当输入包含离散样本点云 $\widehat{\mu}_i = \sum_{j=1}^{m_i} u_{i,j} \delta_{x_{i,j}}$ 时，算法无缝退化为半离散最优传输子问题。对偶变量转化为离散向量 $\bm{\phi} \in \mathbb{R}^{m_i}$，其在连续域 $\Omega$ 上诱导出 Laguerre 胞腔（有界凸多面体）划分： $$L_j(\bm{\phi}) = \left\{ y \in \Omega : \frac{1}{2} \|y - x_{i,j}\|_2^2 - \phi_j \le \frac{1}{2} \|y - x_{i,l}\|_2^2 - \phi_l, \; \forall l \right\}$$ 对偶函数关于分量 $\phi_j$ 的导数简化为理论测度与胞腔积分测度之差 $u_{i,j} - \nu(L_j(\bm{\phi}))$。通过计算该离散势函数的二次 $c$-变换 $\phi^c(y) = \min_{1 \le j \le m_i} \{\frac{1}{2} \|y - x_{i,j}\|_2^2 - \phi_j\}$，直接得到连续重心到离散点集的传输势。对于二维与三维离散点云，通过快速计算 Laguerre 划分，将单步时间复杂度从离散 LP 的 $\mathcal{O}(m_i^3 + M^3)$ 锐降至 $\mathcal{O}(m_i M)$，并在网格对数域通过 $\text{logsumexp}$ 维持极高的数值稳定性。

3. 截断高斯平滑与 $O(\log T / \sqrt{T})$ 次优率收敛保证：兼顾离散逼近与连续紧致

在理论分析层面，当输入全为离散点云时，真实理论重心必然是离散狄拉克测度，使得真实极小值点相对于连续初值 $\rho^0$ 的相对熵发散至无穷大（$D_{\text{KL}}(\lambda^\star \parallel \lambda^0) = +\infty$），直接导致标准凸镜面下降分析失效。为攻坚该退化难题，理论构造了一个参数为 $\varepsilon$ 的截断高斯平滑测度 $\lambda_\varepsilon$ 作为代理桥梁，严密证明了在支撑边界与点云距离满足 $\operatorname{dist}(x_{i,j}, \partial\Omega) \ge r$ 时，不仅 $W_2(\lambda_\varepsilon, \lambda^\star) = \mathcal{O}(\varepsilon)$，且 $D_{\text{KL}}(\rho_\varepsilon \parallel \rho^0) = \mathcal{O}(\log(1/\varepsilon))$。通过平衡离散逼近误差与 KL 散度漂移，在选取步长 $\eta_k = \frac{1}{\sqrt{k+1}}$ 或 $\frac{1}{\sqrt{T}}$ 并设定 $\varepsilon = T^{-1/2}$ 时，严格证明了前 $T$ 步候选重心的目标泛函次优差收敛至： $$\min_{0 \le k \le T} \left( \mathcal{E}(\lambda^k) - \mathcal{E}(\lambda^\star) \right) = \mathcal{O}\left( \frac{\log T}{\sqrt{T}} \right)$$ 这一理论结果在完全连续、完全离散以及离散-连续异构混合输入下全域成立，且仅依赖域 $\Omega$ 的紧凸性，无需以往对偶流方法苛刻的 Sobolev 负范数有界或势函数高阶光滑性假设。

实验关键数据¶

主实验¶

论文在 MNIST 手写数字（784维离散直方图）、异构多模态分布（包含 Swiss roll 点云、心形连续域、MNIST 图像）、2D/3D 高斯点云以及图像色盘平均等多种基准上进行了多维度评测。在 MNIST 数字“3”的重心重构实验中（$n=10$ 张图片，真实真值由大规模线性规划 LP 给出），FRBary 与熵正则化 Wasserstein 重心（EWB）及去偏 Sinkhorn 重心（DSB）的距离误差对比如下：

方法	正则化参数 $\lambda$	$W_2$ 误差（$\times 10^{-4}$）	$L^2$ 密度误差（$\times 10^{-4}$）	运行特性与收敛性
FRBary（本文）	无正则化	1.360	1.573	迭代 500 轮稳定收敛，无熵平滑偏差
EWB (Benamou et al.)	0.0004	1.904	2.875	更小参数（如 0.0003）发生数值溢出崩溃
DSB (Janati et al.)	0.0004	3.237	5.582	去偏调整后边缘处仍呈现非刚性伪影

在异构输入计算耗时对比上（MNIST 图像 + 心形连续密度 + Swiss roll 点云），自由支撑（Free-support）方法与本文方法在相同 250 轮迭代下的耗时（秒）对比如下：

重心计算任务配置	FRBary（本文，$10^4$ 网格连续密度）	Exact Free-support（1000 采样点）	Sinkhorn Free-support（1000 采样点）
左侧两类输入（MNIST + 心形）	275 s	609 s	84 s（早停收敛，点集聚拢于边界）
全部三类输入（图像 + 连续 + 点云）	357 s	798 s	834 s
右侧两类输入（心形 + Swiss roll）	286 s	802 s	468 s

消融实验¶

在不同计算架构与高分辨率直方图求解的性能分析中，论文评估了网格离散化、神经网络参数化与各基线算法在高分辨率手写数字数据集 HWD+（10 张图片）上的求解表现与可扩展性：

算法/配置	分辨率规格	迭代轮数	耗时	空间复杂度	视觉质量与数值状态
FRBary（网格）	$100 \times 100$	200	147 s	$\mathcal{O}(M)$	笔画边缘清晰无模糊
FRBary（网格）	$200 \times 200$	200	423 s	$\mathcal{O}(M)$	平滑保持几何拓扑结构
sGS-ADMM (Yang et al.)	$100 \times 100$	1650	473 s	$\mathcal{O}(M^2)$	笔画正常收敛
sGS-ADMM (Yang et al.)	$200 \times 200$	2950	16,470 s (~4.5 h)	$\mathcal{O}(M^2)$	出现斑块伪影；$250^2$ 爆内存 (>300GB)
WDHA (Kim et al.)	$100 \times 100$	1000	5.5 s	$\mathcal{O}(M)$	依赖高阶 Sobolev 平滑性假设
WDHA (Kim et al.)	$200 \times 200$	1000	67 s	$\mathcal{O}(M)$	无法直接迁移至无规则离散点云
FRBary (神经网络架构)	Swiss Roll ($n=4, 10^4$点)	200	15 min	参数化 MLP	SWD=0.0635，朗之万动力学直接采样
FRBary (半离散网格)	Swiss Roll ($n=4, 10^4$点)	300	11 min	$100 \times 100$ 网格	SWD=0.0272，精度匹敌 Exact-LP (35 min)

关键发现¶

相比传统精确自由支撑线性规划（Exact Free-support），FRBary 无需反复重跑全量 LP 重新估计密度，直接输出连续高分辨率密度，时间开销缩减 55%–64%。
Sinkhorn 自由支撑法虽然在个别任务上提前迭代退出，但其所得解严重退化，往往由于熵正则化失衡而将粒子大量压缩并堆积在支集边界轮廓上。
在样本容量可扩展性测试中，面对 500 张 MNIST 样本构建重心任务，FRBary 依靠势函数并行计算在约 5 分钟内完成，而经典 LP 求解器直接因内存超限（OOM）无法运行。

亮点与洞察¶

无偏几何流设计：巧妙利用负相对熵作为 Bregman 镜像映射，使得算法具有 Fisher-Rao 黎曼流性质，既享受了乘性对数域更新的平滑数值稳定性，又完全避开了 Sinkhorn 算法无法消除的固有熵偏差。
离散点云的连续重构：打破了“离散输入只能求出离散重心点云”的教条，通过半离散传输机制在迭代全程保持连续密度，为分布回归、Wasserstein 主成分分析等依赖连续锚点密度的下游统计任务铺平了道路。
极度轻量且无需精细调参：摆脱了神经网络最优传输对庞大判别器网络、稳定训练技巧及极大极小鞍点调优的脆弱依赖，在 2D/3D 上仅需一阶凸优化步长衰减规则即可达到极高几何保真度。

局限与展望¶

维度扩展受阻：目前基于网格划分与经典半离散计算库（如 sdot、pysdot）的高精度实现主要局限于 2 维和 3 维问题。面对更高维场景时，网格单元数随维度指数爆炸。
高维采样偏差：虽然论文给出了结合 MLP 与朗之万动力学的连续神经网络变体，但在 3 维以上空间中，简单未调整朗之万动力学（ULD）会引入不可忽视的采样偏差，而切换至 HMC 或 NUTS 采样器会导致计算时间难以接受。
非紧凸域拓展：目前的收敛理论严格依赖底层支撑集 $\Omega$ 的紧致性与凸性假设，未来值得进一步向无界欧氏全空间或具有黎曼曲率的紧流形空间推广。

评分¶

新颖性: ⭐⭐⭐⭐☆ 创造性地结合 Fisher-Rao 镜面下降与半离散传输，完美统一离散与连续测度重心求解。
实验充分度: ⭐⭐⭐⭐☆ 覆盖点云、手写字图像、连续几何域及异构输入，对比了 ADMM、Sinkhorn、神经网络等完备基线。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严密完整，从连续推导平滑过渡至离散高斯截断逼近，结构一气呵成。
价值: ⭐⭐⭐⭐☆ 为计算几何、统计最优化与图像色调迁移等领域提供了高可靠性、无偏且结构优雅的开源求解基线。

算法/配置	分辨率规格	迭代轮数	耗时	空间复杂度	视觉质量与数值状态
FRBary（网格）	\(100 \times 100\)	200	147 s	\(\mathcal{O}(M)\)	笔画边缘清晰无模糊
FRBary（网格）	\(200 \times 200\)	200	423 s	\(\mathcal{O}(M)\)	平滑保持几何拓扑结构
sGS-ADMM (Yang et al.)	\(100 \times 100\)	1650	473 s	\(\mathcal{O}(M^2)\)	笔画正常收敛
sGS-ADMM (Yang et al.)	\(200 \times 200\)	2950	16,470 s (~4.5 h)	\(\mathcal{O}(M^2)\)	出现斑块伪影；\(250^2\) 爆内存 (>300GB)
WDHA (Kim et al.)	\(100 \times 100\)	1000	5.5 s	\(\mathcal{O}(M)\)	依赖高阶 Sobolev 平滑性假设
WDHA (Kim et al.)	\(200 \times 200\)	1000	67 s	\(\mathcal{O}(M)\)	无法直接迁移至无规则离散点云
FRBary (神经网络架构)	Swiss Roll (\(n=4, 10^4\)点)	200	15 min	参数化 MLP	SWD=0.0635，朗之万动力学直接采样
FRBary (半离散网格)	Swiss Roll (\(n=4, 10^4\)点)	300	11 min	\(100 \times 100\) 网格	SWD=0.0272，精度匹敌 Exact-LP (35 min)