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

技术摘要：Wasserstein 空间中的差分凸规划及其在 MMD 优化中的应用

问题陈述
在概率测度空间上优化泛函是机器学习任务的核心，涵盖了从变分推断到生成模型等广泛领域。一种标准方法是通过离散化 Wasserstein 梯度流，直接在 Wasserstein 空间 $(P_2(\mathbb{R}^d), W_2)$ 上进行优化。虽然对于凸泛函（如 KL 散度、f-散度）这种方法非常有效，但许多具有实际意义的目标函数在 Wasserstein 测地线上是非凸的。显著的例子包括平方 Wasserstein 距离、切片 Wasserstein 距离以及最大均值差异（MMD）。标准的阶一方法，如 Wasserstein 梯度下降（WGD），在这些非凸设置中表现不佳，通常需要精细的噪声注入，或者在替代空间（例如 Wasserstein-MMD 或 Wasserstein-Fisher-Rao 空间）中进行优化，而后者会改变分布权重或引入隐式步骤。

方法论
作者提出了一种针对具有 差分凸（Difference-of-Convex, DC） 分解性质的非凸泛函量身定制的新型粒子算法。其核心方法是将经典的凸-凹程序（Convex-Concave Procedure, CCCP）从欧几里得空间提升到 Wasserstein 空间，从而得到了 Wasserstein 凸-凹程序（WCCCP）。

1. Wasserstein 空间中的 DC 分解： 目标泛函 $F$ 被分解为 $F(\mu) = F_+(\mu) - F_-(\mu)$，其中 $F_+$ 和 $F_-$ 都是 $P_2(\mathbb{R}^d)$ 上的全凸（totally convex）泛函。与以往将凹部分限制在势能（线性泛函）上的工作不同，本方法允许 $F_-$ 为任何可微泛函。
2. 算法方案： 在每次迭代 $k$，给定测度 $\mu_k$，算法求解一个极大化-极小化（majorization-minimization）问题。它在 $\mu_k$ 附近对凹部分 $F_-$ 进行线性化，以构造一个上界，并在传输映射空间 $T \in L^2(\mu_k)$ 上最小化所得的凸代理函数：
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   接下来的迭代通过定义 $\mu_{k+1} = (T_{k+1})\#\mu_k$ 来实现。
3. 实现： 内部最小化过程通过在 $L^2(\mu_k)$ 上的梯度下降来完成。对于经验分布（粒子），这转化为更新粒子位置。作者还建立了 WCCCP 与 Wasserstein 空间下的镜像下降（Mirror Descent）及 Bregman 近端下降（Bregman Proximal Descent）之间的联系。
4. MMD 特性： 对于最大均值差异（MMD），作者通过将核函数 $k(x,y) = \psi(x-y)$ 本身分解为凸分量 $\psi = \psi_+ - \psi_-$，提出了显式的 DC 分解。他们为平移不变核（包括高斯核和 Riesz 核）提供了具体的分解方法，利用 Jordan 分解或代数拆分（例如 $\cosh/\sinh$）。

核心贡献

• WCCCP 算法： 引入了 Wasserstein 凸-凹程序，将 DC 优化扩展到 Wasserstein 空间，且不限制凹分量必须为线性势能。
• 理论收敛性：
  • 证明了沿迭代路径目标函数值 $F(\mu_k)$ 是非增的。
  • 建立了**几乎平稳性（almost stationarity）**结果。在组件满足光滑性和强凸性（全凸性）的假设下，作者证明了迭代间距离平方与 Wasserstein 梯度范数之间的亚线性收敛率。具体而言，$\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$。
  • 论证了该方案在特定条件下等价于 Bregman 近端下降和镜像下降。
• MMD 的显式 DC 分解： 本文提供了一种系统的方法，将具有平移不变核的 MMD 泛函分解为满足必要凸性和光滑性条件的 DC 形式。
• 实证验证： 在能量距离（Energy Distance）和带有高斯核的 MMD 实验中，WCCCP 比标准的 WGD 和 Wasserstein 近端梯度（前向-后向）方法表现出更快的收敛速度和更高的稳定性。

结果

• 能量距离： 在合成数据（螺旋形和猫形）以及 CIFAR10 图像上，WCCCP 的收敛速度明显快于 WGD 和前向-后向（FB）方案。该算法成功通过最小化能量距离来匹配复杂的特征分布。
• 高斯核 MMD： 最小化具有高斯核的 MMD 对带宽和初始化极其敏感。作者展示了 WGD 和使用标准分解的 FB 方案往往会陷入局部极小值，导致粒子偏离目标。相比之下，使用特定核分解（Jordan 或 $\cosh/\sinh$）的 WCCCP 实现了更好的收敛。研究指出，Jordan 分解更加光滑且不易产生数值不稳定，而 $\cosh/\sinh$ 分解则确保所有粒子最终都能收敛，尽管速率较慢。
• 对比： 结果强调了 WCCCP 的性能高度依赖于 DC 分解的选择。一个“经过良好选择”的分解对于超越标准梯度下降至关重要。

意义与主张
本文声称提供了一个在 Wasserstein 空间中优化非凸泛函的鲁棒框架，解决了标准梯度方法因非凸性而失效的问题。通过将 CCCP 提升到此设置，作者提供了一个具有理论依据的替代方案，避免了以往针对 MMD 方法中存在的权重混合问题。这项工作强调，虽然算法具有通用性，但其实际成功很大程度上取决于核函数的特定 DC 分解。作者谦虚地指出，未来仍需进一步研究以更深入地理解不同分解的影响，并设计自动化的自适应策略来选择分解方式。此外，目前的理论分析假设内部优化问题具有精确解，将如何处理不精确求解器确定为未来的研究方向。