Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization¶

会议: NeurIPS 2026
arXiv: 2606.27767
OpenReview: https://openreview.net/forum?id=xKLfU6Cqa2
代码: https://github.com/clbonet/Wasserstein_Convex_Concave_Procedure
领域: 优化/理论
关键词: Wasserstein空间, 凸差规划, 凸凹过程, 最大均值差异, 能量距离

一句话总结¶

针对概率测度空间中最大均值差异（MMD）等目标沿测地线非凸导致的优化停滞问题，提出了 Wasserstein 凸凹过程（WCCCP），利用平移不变核与径向核的凸差分解实现亚线性收敛与稳定粒子优化。

研究背景与动机¶

在机器学习中，对概率测度空间上的泛函进行最优化已广泛应用于变分推断、生成建模、强化学习、神经网络均场优化及细胞动力学建模。主流方法是将概率分布空间赋予 \(W_2\) Wasserstein 距离，并通过时间离散化求解对应的 Wasserstein 梯度流（Wasserstein Gradient Flow）。这衍生出了诸如 Wasserstein 梯度下降（WGD）、近端点与近端梯度算法、坐标下降和镜像下降等一系列欧氏优化算法在测度空间上的对应物。

然而，现有的收敛性分析大多强依赖于泛函沿广义测地线的凸性（geodesic convexity）。实践中极具价值的多种目标泛函在 Wasserstein 空间中本质上都是非凸的，例如最大均值差异（Maximum Mean Discrepancy, MMD）、能量距离（Energy Distance, ED）、非对数凹目标分布下的 KL 散度以及切片 Wasserstein 距离。虽然部分学者尝试通过注入高斯噪声或在权重可变的空间（如 Wasserstein-Fisher-Rao 或 Wasserstein-MMD 空间）中做松弛优化，但噪声尺度调优极为敏感，且修改粒子权重改变了标准粒子传输的几何特性。

欧氏空间中解决非凸问题的经典利器是凸差规划（Difference-of-Convex, DC）与凸凹过程（Convex-Concave Procedure, CCCP），但将其推广至测度空间面临严峻挑战：既有工作仅考虑了凹项为线性势能的特殊情况，无法处理 MMD 这种包含非凸二次交互能量项的双重非凸结构。核心 idea：将经典凸凹过程拓展到 Wasserstein 空间提出 WCCCP，并基于核函数本身的代数结构与 Jordan 积分分解，构造出 MMD 与能量距离在 Wasserstein 空间中的显式 DC 分解，彻底消除粒子传输优化中的伪局部极值与发散问题。

方法详解¶

整体框架¶

WCCCP 旨在最小化形式为 \(\mathcal{F}(\mu) = \mathcal{F}^+(\mu) - \mathcal{F}^-(\mu)\) 的目标泛函，其中 \(\mathcal{F}^+\) 与 \(\mathcal{F}^-\) 均为 \(\mathcal{P}_2(\mathbb{R}^d)\) 上的完全凸（totally convex）泛函。算法在每轮外层迭代中，利用 \(\mathcal{F}^-\) 的完全凸性建立线性化下界，从而构造关于当前测度推前映射 \(T\) 的凸上界并予以极小化。

对于带有平移不变核 \(k(x, y) = \psi(x - y)\) 的 MMD，若核母函数具备标量 DC 分解 \(\psi = \psi^+ - \psi^-\)，则目标中的交互能量与势能项被解耦重组为两个完全凸泛函的差。外层通过粒子推前更新测度分布，内层子问题则在切空间 \(L^2(\mu_k)\) 上利用梯度下降精确求解最优传输映射。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：源粒子分布 μ_0 与目标经验测度 ν"] --> B["核函数 DC 构造<br/>基于径向 Jordan 分解或代数分裂得到 ψ+ 与 ψ-"]
    B --> C["泛函分裂<br/>将 MMD 解耦为完全凸项 F+ 与 F-"]
    C --> D["WCCCP 外层主循环<br/>线性化凹项 -F- 构造 Bregman 替代上界"]
    D --> E["内层粒子求解器<br/>在 L2 空间沿 ∇W2F+ 与 ∇W2F- 梯度下降求传输映射"]
    E --> F["测度推前更新<br/>μ_k+1 = (T_k+1)# μ_k 驱动粒子流向目标"]
    F -->|判据 D_F+^k > 0| D
    F -->|达到准平稳点| G["输出：收敛粒子分布 μ* ≈ ν"]

关键设计¶

1. Wasserstein 凸凹过程：泛函线性化上界与等价镜像解释

传统一阶 Wasserstein 算法在遇到非凸能面时极易停滞。WCCCP 在第 \(k\) 步将凹项 \(-\mathcal{F}^-(\mathrm{T}_{\#}\mu_k)\) 围绕恒等映射 \(\mathrm{Id}\) 线性化。由于 \(\mathcal{F}^-\) 是完全凸泛函，由 Bregman 散度 \(\mathrm{D}_{\mathcal{F}^-}^{\mu_k}(\mathrm{T}, \mathrm{Id}) \geq 0\) 可得全局支撑切平面，进而推导得到 \(\mathcal{F}(\mathrm{T}_{\#}\mu_k)\) 的严格代理上界：

\[\mathcal{F}(\mathrm{T}_{\#}\mu_k) \leq \mathcal{F}^+(\mathrm{T}_{\#}\mu_k) - \mathcal{F}^-(\mu_k) - \langle \nabla_{\mathrm{W}_2}\mathcal{F}^-(\mu_k), \mathrm{T} - \mathrm{Id} \rangle_{L^2(\mu_k)}\]

极小化该上界确定下一步传输映射 \(\mathrm{T}_{k+1} = \operatorname{argmin}_{\mathrm{T} \in L^2(\mu_k)} \mathcal{F}^+(\mathrm{T}_{\#}\mu_k) - \langle \nabla_{\mathrm{W}_2}\mathcal{F}^-(\mu_k), \mathrm{T} - \mathrm{Id} \rangle_{L^2(\mu_k)}\)，随后执行推前测度更新 \(\mu_{k+1} = (\mathrm{T}_{k+1})_{\#}\mu_k\)。作者证明，该过程在步长 \(\tau=1\) 下同时严格等价于以 \(\mathcal{F}^+\) 为势能的 Wasserstein 镜像下降，以及以 \(\mathcal{F}^-\) 为势能的 Wasserstein Bregman 近端下降。当 \(\mathcal{F}^+\) 具 \(W_2\) 可微性时，外层一阶驻点最优性条件满足 \(\nabla_{\mathrm{W}_2}\mathcal{F}^+(\mu_{k+1}) \circ \mathrm{T}_{k+1} = \nabla_{\mathrm{W}_2}\mathcal{F}^-(\mu_k)\)。

2. 核函数 DC 分解：平移不变核与径向 Jordan 分裂

MMD 的二次交互能 \(\iint k(x,y)\mathrm{d}\mu(x)\mathrm{d}\mu(y)\) 与交叉势能 \(-\int k(x,y)\mathrm{d}\nu(y)\) 互相耦合且沿测地线非凸。直接引入外加二次正则化（如 Luu 等人做法）不仅扭曲了能量曲率，且要求极大的正则系数。本文提出直接分解平移不变核 \(k(x, y) = \psi(x - y) = \psi^+(x - y) - \psi^-(x - y)\)。此时 MMD 自然拆解为两个完全凸泛函之差：

\[\mathcal{F}^+(\mu) = \frac{1}{2}\iint \psi^+(x-y)\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int V^-\mathrm{d}\mu + c(\nu), \quad V^-(\cdot) = \int \psi^-(\cdot - y)\mathrm{d}\nu(y)\]

\[\mathcal{F}^-(\mu) = \frac{1}{2}\iint \psi^-(x-y)\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int V^+\mathrm{d}\mu, \quad V^+(\cdot) = \int \psi^+(\cdot - y)\mathrm{d}\nu(y)\]

针对径向核 \(\psi(z) = q(\|z\|_2^2)\)，为保证生成的 \(\psi^\pm\) 在紧凸集上强凸，作者推导了径向 Hessian 特征值界，并给出通用的 Jordan 积分构造法：

\[q_+(s) = q(0) + (q'(0) + A)s + \int_0^s (s-t)\max(0, q''(t))\mathrm{d}t, \quad q_-(s) = As - \int_0^s (s-t)\min(0, q''(t))\mathrm{d}t\]

其中 \(A \geq \max(0, -q'(0))\)。特别地，对于高斯核 \(q(s) = e^{-s/(2h)}\)，选取 \(A = 1/(2h)\) 导出光滑的 Gauss-Jordan 分解 \(q_+(s) = e^{-s/(2h)} + s/(2h)\) 和 \(q_-(s) = s/(2h)\)；或利用解析展开将偶数与奇数阶项拆解为双曲余弦与双曲正弦 Gauss-cosh/sinh 分解：\(q_+(s) = \cosh(s/(2h))\)，\(q_-(s) = \sinh(s/(2h))\)。

3. 内层粒子梯度求解器：在离散经验分布上的时间离散化迭代

在实际粒子系统中，测度由 \(n\) 个离散粒子经验分布表示 \(\mu_k = \frac{1}{n}\sum_{i=1}^n \delta_{x_i^{k,0}}\)。虽然内层子问题通常不存在显式闭式解，但由于子问题目标泛函在切空间 \(L^2(\mu_k)\) 上是严格凸的，且其梯度具有 \(L\)-Lipschitz 连续性（\(L = \sqrt{2}\overline{\Lambda}[q_+] + \overline{\Lambda}[q_-]\)），可直接运用内层有限步梯度下降（或带动量的梯度下降）：

\[x_i^{k, \ell+1} = x_i^{k, \ell} - \tau \left( \nabla_{\mathrm{W}_2}\mathcal{F}^+(\mu_k^\ell)(x_i^{k, \ell}) - \nabla_{\mathrm{W}_2}\mathcal{F}^-(\mu_k)(x_i^{k, 0}) \right)\]

经过 \(M\) 步内层更新后，令 \(x_i^{k+1, 0} = x_i^{k, M}\) 完成一轮外层步。这种粒子级两层解算机制完全避免了外加扩散噪声破坏分布支撑集的问题，同时在大规模粒子流动中保持了强数值稳定性。

损失函数 / 训练策略¶

在外层收敛性监控中，引入代理残差准则 \(\mathcal{D}_{\mathcal{F}^+}^k \coloneqq \mathcal{F}^+(\mu_k) - \mathcal{F}^+(\mu_{k+1}) - \langle \nabla_{\mathrm{W}_2}\mathcal{F}^-(\mu_k), \mathrm{Id} - \mathrm{T}_{k+1} \rangle_{L^2(\mu_k)}\)。当 \(\mathcal{F}^+\) 可微时，\(\mathcal{D}_{\mathcal{F}^+}^k = \mathrm{D}_{\mathcal{F}^+}^{\mu_k}(\mathrm{Id}, \mathrm{T}_{k+1}) \geq 0\)。

理论证明，在泛函有下界且具有强完全凸性（\(\alpha^+ + \alpha^- > 0\)）与 \(L\)-Lipschitz 梯度条件下，WCCCP 序列满足亚线性收敛率：

\[\min_{0 \leq k \leq K-1} \|\nabla_{\mathrm{W}_2}\mathcal{F}(\mu_k)\|_{L^2(\mu_k)}^2 \leq \frac{2L^2}{\alpha^+ + \alpha^-} \frac{\mathcal{F}(\mu_0) - \mathcal{F}(\mu_K)}{K} = \mathcal{O}\left(\frac{1}{K}\right)\]

即算法在最多 \(K\) 步外层迭代中必定遍历到至少一个 \(\mathcal{O}(1/\sqrt{K})\) 的近驻点。

实验关键数据¶

主实验¶

在低维流形形状（Spiral、Cat、Heart、Disk）与真实高维图像数据集（CIFAR-10, \(d=3072\)；MNIST, \(d=784\)）上，以 \(n=500\) 粒子从各向同性高斯初始化 \(\mathcal{N}(0, I_d)\) 出发，分别评估能量距离（ED）与带核带宽 \(h=10\) 的高斯 MMD 优化表现。对比基准包括标准 Wasserstein 梯度下降（WGD）与基于外加势能正则化的 Wasserstein 近端梯度（Forward-Backward, FB）。

任务 / 数据集	目标函数 / 核配置	方法	外层/内层步数	计算用时	收敛表现 / 终止目标值
CIFAR-10 (\(d=3072\))	能量距离 (Negative Distance)	WGD (\(\tau=1\))	200,000 / 0	1h30m (V100)	陷入平台期，远离收敛
CIFAR-10 (\(d=3072\))	能量距离 (Negative Distance)	WCCCP (Ours)	40,000 / 20	1h15m (V100)	完全收敛，精准复现目标支撑集
MNIST (\(d=784\))	能量距离 (Negative Distance)	WGD (\(\tau=1\))	200,000 / 0	~18m (V100)	缓慢收敛，边缘伪影较重
MNIST (\(d=784\))	能量距离 (Negative Distance)	WCCCP (Ours)	40,000 / 20	~18m (V100)	快速平稳收敛至目标经验分布
2D 高斯混合 (3 模态)	高斯 MMD (\(h=10\))	WGD (\(\tau=1\))	200,000 / 0	-	部分粒子被弹射发散，严重陷入局部极值
2D 高斯混合 (3 模态)	高斯 MMD (\(h=10\))	FB (势能二次正则)	200,000 / 250	-	停滞于远离真实模式的局部最优
2D 高斯混合 (3 模态)	高斯 MMD (\(h=10\))	WCCCP (Gauss-Jordan)	200,000 / 500	-	前期迅速下降，无粒子发散
2D 高斯混合 (3 模态)	高斯 MMD (\(h=10\))	WCCCP (cosh/sinh)	200,000 / 250	-	粒子 100% 收敛至三模态支撑域

消融实验¶

研究高斯核 MMD 优化中不同 DC 分解机制对优化轨迹与数值稳定性的影响（\(n=500\) 粒子，源分布 \(\mathcal{N}(5\mathbb{1}_2, I_2)\)，单高斯目标分布 \(\mathcal{N}(0, \Sigma)\)，25 次独立重复实验统计）：

DC 分解方案	凹凸分解形式 \(q_+(s) / q_-(s)\)	数值稳定性	粒子发散率	收敛速率与质量
既有基线 (Luu et al.)	外加全域势能 \(L\\|x\\|^2_2\) 正则化	极差	高（大量粒子飞出视界）	与 WGD 类似，困于局部极小
Gauss-Jordan (Ours)	\(e^{-s/(2h)} + \frac{s}{2h} \;/\; \frac{s}{2h}\)	优异（无指数爆炸）	极低（极少陷入较浅极小）	前期下降最快，步长鲁棒（\(\tau=0.1\)）
Gauss-cosh/sinh (Ours)	\(\cosh(\frac{s}{2h}) \;/\; \sinh(\frac{s}{2h})\)	需小步长防梯度溢出	0%（全部粒子均正确汇聚）	最终误差最低，单调稳定收敛

关键发现¶

针对高斯核 MMD 经典存在的“粒子发散”（particles drifting away）与维数灾难缺陷，WCCCP 搭配基于核本身的径向分解（Gauss-cosh/sinh 或 Gauss-Jordan）彻底根治了模式排斥问题，而既有文献提出的外加势能二次分解与 WGD 均无法避免粒子发散。
Gauss-Jordan 分解由于其各阶导数受控、没有双曲函数随距离增加呈指数级恶性膨胀的问题，可以使用相对较大的内层学习率（\(\tau=0.1\) 对比 cosh/sinh 的 \(5\times 10^{-4}\)），在计算成本受控的情况下收敛最为迅捷。
在 CIFAR-10 高维（\(d=3072\)）图像实验中，在严格匹配总运算时间（WCCCP 1小时15分 vs WGD 1小时30分）的公平前提下，WCCCP 的收敛速度远超传统 Wasserstein 梯度下降，显示出其在生成式粒子建模领域的广阔应用前景。

亮点与洞察¶

从点空间到测度空间的理论升维：成功将非凸欧氏优化中的核心支柱 CCCP 严格推广至 Wasserstein 空间，建立了基于 Bregman 散度与完全凸性（Total Convexity）的完整收敛性理论体系，证明了 \(\mathcal{O}(1/K)\) 的近驻点收敛界。
核内生分解替代粗暴正则：摒弃了以往对测度整体强加二次势能这一导致病态曲率的粗糙做法，创造性地直接通过对平移不变核及径向二阶导数的 Jordan 积分进行内生分裂，保持了能量泛函天然的物理与几何对称性。
跨范式等价性透视：深刻揭示了 WCCCP 与 Wasserstein 镜像下降（Mirror Descent）以及 Wasserstein Bregman 近端下降的本征等价性，打通了测度空间中非凸 DC 算法与现代分裂算法之间的内在纽带。

局限与展望¶

依赖预设的核函数分解：当前核分解方案主要依赖于作者针对径向平移不变核手工设计的代数或积分形式，尚未结合优化几何自适应、自动化地寻找在当前点曲率最优的分解配对。
内层子问题不精确解的理论缺口：理论收敛率推导建立在内层极小化问题能够精确求解的前提下，而实际算法采用有限步梯度下降近似，不精确解（inexact solver）对全局收敛阶数与渐进驻点的扰动分析仍有待深入完善。
Wasserstein-PL 不等式尚未完全闭合：尽管探讨了连续测度空间中的 DC-PL 不等式，但由于 Wasserstein 空间上尚缺乏普适完备的 Fenchel 对偶工具体系，尚无法建立如同欧氏空间般无条件的全局线性收敛证明。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性建立 Wasserstein 空间下的 DC 规划与 CCCP 理论，并给出径向核分解框架。
实验充分度: ⭐⭐⭐⭐☆ 覆盖了低维拓扑形状、多模态高斯混合以及高维真实图像生成流（CIFAR-10/MNIST），基线对齐公平。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，从泛函分析到离散粒子算法层层递进，结构规范明晰。
价值: ⭐⭐⭐⭐⭐ 为利用测度梯度流解决统计检验、生成对抗拟合与非凸均场动力学提供了全新强力工具。