Wasserstein空间中的凸差规划及其在MMD优化中的应用

# Wasserstein空间中的凸规划差分及其在MMD优化中的应用  
**来源**：https://arxiv.org/html/2606.27767  

**Clément Bonet**  
CMAP, CNRS, 巴黎综合理工学院, IP Paris  
clemont\.bonet\.mapp@polytechnique\.edu  

**Pierre-Cyril Aubin-Frankowski**  
CERMICS, CNRS, 巴黎桥路学院, IP Paris  
pierre-cyril\.aubin@enpc\.fr  

**Youssef Mroueh**  
IBM Research  
mroueh@us\.ibm\.com  

###### 摘要

在概率测度空间上优化泛函已成为机器学习中的常见问题。一种广泛使用的方法是直接在Wasserstein空间上进行优化，但许多实际目标泛函沿Wasserstein测地线是非凸的，这使得对标准一阶方法的分析变得困难。在本工作中，我们研究了一类在Wasserstein空间上具有差凸（DC）分解的目标函数，并将经典的凸凹过程（CCCP）推广到该设定下。在分解的凸分量满足光滑性和强凸性的假设下，我们证明了所提算法迭代点的几乎平稳性。我们的主要关注点是最大均值差异（MMD）和能量距离（ED）泛函，我们为它们建立了显式的Wasserstein DC分解，并在温和假设下证明了该方案的局部收敛性。实验表明，对于这些MMD目标，精心选择的DC分解能比Wasserstein梯度下降实现更快、更稳定的收敛。

### 1 引言

在概率测度空间上优化是机器学习中的一个重要问题，已广泛应用于变分推断（Blei等人，2017 (https://arxiv.org/html/2606.27767#bib.bib42)；Lambert等人，2022 (https://arxiv.org/html/2606.27767#bib.bib52)；Petit-Talamon等人，2025 (https://arxiv.org/html/2606.27767#bib.bib53)）、生成建模（Deng等人，2026 (https://arxiv.org/html/2606.27767#bib.bib38)；Turan和Ovsjanikov，2026 (https://arxiv.org/html/2606.27767#bib.bib39)；Cao等人，2026 (https://arxiv.org/html/2606.27767#bib.bib40)）、强化学习（Zhang等人，2018 (https://arxiv.org/html/2606.27767#bib.bib54)；Pfau等人，2025 (https://arxiv.org/html/2606.27767#bib.bib55)）、神经网络优化（Mei等人，2018 (https://arxiv.org/html/2606.27767#bib.bib72)；Chizat和Bach，2018 (https://arxiv.org/html/2606.27767#bib.bib73)）以及细胞动力学建模（Bunne等人，2022 (https://arxiv.org/html/2606.27767#bib.bib75)；Terpin等人，2024 (https://arxiv.org/html/2606.27767#bib.bib78)；Persiianov等人，2026 (https://arxiv.org/html/2606.27767#bib.bib77)）等问题。在概率测度空间上优化的一个主要方式是赋予其Wasserstein距离（Villani等人，2009 (https://arxiv.org/html/2606.27767#bib.bib79)；Santambrogio，2015 (https://arxiv.org/html/2606.27767#bib.bib51)），并离散化相应的Wasserstein梯度流（Jordan等人，1998 (https://arxiv.org/html/2606.27767#bib.bib44)；Wibisono，2018 (https://arxiv.org/html/2606.27767#bib.bib45)）。这使得我们可以设计许多与欧几里得版本相对应的优化算法，例如梯度下降（Wibisono，2018 (https://arxiv.org/html/2606.27767#bib.bib45)）、近端点与梯度算法（Jordan等人，1998 (https://arxiv.org/html/2606.27767#bib.bib44)；Salim等人，2020 (https://arxiv.org/html/2606.27767#bib.bib12)）、坐标下降（Xu和Li，2026 (https://arxiv.org/html/2606.27767#bib.bib25)）或镜像下降（Bonet等人，2024 (https://arxiv.org/html/2606.27767#bib.bib3)；Sharrock等人，2023 (https://arxiv.org/html/2606.27767#bib.bib43)）。尽管这些优化方法在最小化诸如KL散度（Wibisono，2018 (https://arxiv.org/html/2606.27767#bib.bib45)）、ff-散度（Gao等人，2019 (https://arxiv.org/html/2606.27767#bib.bib63)；Ansari等人，2021 (https://arxiv.org/html/2606.27767#bib.bib62)；Liu等人，2024 (https://arxiv.org/html/2606.27767#bib.bib64)）、能量距离（Hertrich等人，2024b (https://arxiv.org/html/2606.27767#bib.bib27)）或切片Wasserstein距离（Liutkus等人，2019 (https://arxiv.org/html/2606.27767#bib.bib46)；Du等人，2023 (https://arxiv.org/html/2606.27767#bib.bib47)；Bonet等人，2025 (https://arxiv.org/html/2606.27767#bib.bib48)）等目标方面取得了良好效果，但其中大多数方法都是针对沿（广义）测地线凸的泛函而设计的。然而，许多常用的泛函是已知非凸的。例如，平方Wasserstein距离本身（Ambrosio等人，2008 (https://arxiv.org/html/2606.27767#bib.bib31)，第9章）、切片Wasserstein距离（Bonnotte，2013 (https://arxiv.org/html/2606.27767#bib.bib50)）、针对非对数凹目标的KL散度（Luu等人，2024 (https://arxiv.org/html/2606.27767#bib.bib2)）、能量距离（Hertrich等人，2024a (https://arxiv.org/html/2606.27767#bib.bib26)）或最大均值差异（MMD）（Arbel等人，2019 (https://arxiv.org/html/2606.27767#bib.bib22)）等都属于此类。凸性假设的一个替代方案是考虑Polyak-Łojasiewicz不等式（Blanchet和Bolte，2018 (https://arxiv.org/html/2606.27767#bib.bib65)；Liu等人，2023 (https://arxiv.org/html/2606.27767#bib.bib66)；Zhu和Chen，2025 (https://arxiv.org/html/2606.27767#bib.bib74)），但迄今这些不等式仅在限制性条件下成立，例如平方Wasserstein距离的平凡情况、满足对数Sobolev不等式测度下的KL散度（Villani等人，2009 (https://arxiv.org/html/2606.27767#bib.bib79)，第21章）、高斯分布上的切片Wasserstein距离（Thurin等人，2026 (https://arxiv.org/html/2606.27767#bib.bib49)）以及带库仑核与光滑初始化的MMD（Boufadène和Vialard，2025 (https://arxiv.org/html/2606.27767#bib.bib100)，第2节）。针对MMD在Wasserstein空间中的非凸性，Arbel等人（2019 (https://arxiv.org/html/2606.27767#bib.bib22)）提出了向梯度注入噪声的方法。然而，噪声量的调节仍然很棘手。Gladin等人（2024 (https://arxiv.org/html/2606.27767#bib.bib24)）则转而将MMD优化置于Wasserstein-MMD空间中（其中MMD是凸的）。他们的方法提升了性能，但需要改变分布的权重，混合了平方Wasserstein和MMD的两步隐式步骤。最近，Belhadji等人（2026 (https://arxiv.org/html/2606.27767#bib.bib106)）通过不动点算法在Wasserstein Fisher-Rao空间中进行优化，同样改变了权重。我们则提出一种新的基于粒子的Wasserstein空间算法来处理非凸泛函。为此，Rd\mathbb{R}^d上的一类类比算法是针对可写为差凸函数（DC）的目标函数的方法族（Le Thi和Pham Dinh，2018 (https://arxiv.org/html/2606.27767#bib.bib70)；Pham Dinh和Le Thi，2014 (https://arxiv.org/html/2606.27767#bib.bib69)）。在Rd\mathbb{R}^d上，这包括了大量感兴趣的函数，特别是两次连续可微函数（Yuille和Rangarajan，2001 (https://arxiv.org/html/2606.27767#bib.bib8)；Hiriart-Urruty，1985 (https://arxiv.org/html/2606.27767#bib.bib94)）。其中一种关键的算法是凸凹过程（Yuille和Rangarajan，2001 (https://arxiv.org/html/2606.27767#bib.bib8)）。这些DC算法已用于机器学习中的许多应用，例如核选择（Argyriou等人，2006 (https://arxiv.org/html/2606.27767#bib.bib56)）、聚类（Tao等人，2014 (https://arxiv.org/html/2606.27767#bib.bib82)）、字典学习（Vo等人，2015 (https://arxiv.org/html/2606.27767#bib.bib83)）、最优传输（Tran等人，2021 (https://arxiv.org/html/2606.27767#bib.bib86)）或神经网络优化（Awasthi等人，2024 (https://arxiv.org/html/2606.27767#bib.bib84)；Askarizadeh等人，2024 (https://arxiv.org/html/2606.27767#bib.bib85)）等。然而，除了少数研究在黎曼流形上考虑DC算法（Souza和Oliveira，2015 (https://arxiv.org/html/2606.27767#bib.bib89)；Weber和Sra，2023 (https://arxiv.org/html/2606.27767#bib.bib17)；Bergmann等人，2024 (https://arxiv.org/html/2606.27767#bib.bib87)；Ferreira等人，2026 (https://arxiv.org/html/2606.27767#bib.bib88)），以及最近在Wasserstein空间上的工作（Luu等人，2024 (https://arxiv.org/html/2606.27767#bib.bib2)；Luu和Wang，2026 (https://arxiv.org/html/2606.27767#bib.bib104)）外，DC泛函的优化算法大多是在欧几里得空间上研究的。

##### 贡献。  
在本工作中，我们专注于在Wasserstein空间上针对可写为凸目标之差的非凸目标开发优化方案。为此，我们引入了*Wasserstein凸凹过程*（WCCCP），并分析了其理论收敛性，证明了其迭代点几乎平稳。与本文最相关的工作是Luu等人（2024 (https://arxiv.org/html/2606.27767#bib.bib2)），他们使用了Wasserstein近端梯度算法（Salim等人，2020 (https://arxiv.org/html/2606.27767#bib.bib12)）来解决某些DC问题，但将应用限制在其凹部为势能（即线性）的泛函。我们认为这对于处理像最大均值差异这样的目标过于局限，因为MMD可以分解为两个非凸二次项与线性项之和。因此，我们处理更一般的情形，其中凹部可以是Wasserstein空间上的任意可微泛函。然后，对于几种核，我们证明了通过适当拆分核得到分解后，所提方案能够比Wasserstein梯度下降更好地优化MMD。特别地，我们提供了关于能量距离和高斯核MMD的实验。

##### 符号说明。  
我们用P2(Rd)\mathcal{P}_{2}(\mathbb{R}^{d})表示具有有限二阶矩的概率分布空间，用Pac(Rd)\mathcal{P}_{\mathrm{ac}}(\mathbb{R}^{d})表示其在Lebesgue测度下绝对连续测度的子集。给定μ,ν∈P2(Rd)\mu,\nu\in\mathcal{P}_{2}(\mathbb{R}^{d})，我们用W22(μ,ν)=infγ∈Π(μ,ν)∫‖x−y‖22dγ(x,y)\mathrm{W}_{2}^{2}(\mu,\nu)=\inf_{\gamma\in\Pi(\mu,\nu)}\int\|x-y\|_{2}^{2}\ \mathrm{d}\gamma(x,y)表示平方Wasserstein距离，其中Π(μ,ν)\Pi(\mu,\nu)是μ\mu与νν之间的耦合集，Πo(μ,ν)\Pi_{o}(\mu,\nu)是其最优耦合子集。度量空间(P2(Rd),W2)(\mathcal{P}_{2}(\mathbb{R}^{d}),\mathrm{W}_{2})称为Wasserstein空间。对任意μ∈P2(Rd)\mu\in\mathcal{P}_{2}(\mathbb{R}^{d})，我们用L2(μ)L^{2}(\mu)表示满足∫‖f(x)‖22dμ(x)<∞\int\|f(x)\|_{2}^{2}\ \mathrm{d}\mu(x)<\infty的函数f:Rd→Rdf:\mathbb{R}^{d}\to\mathbb{R}^{d}的Hilbert空间，配备范数‖f‖L2(μ)2=∫‖f(x)‖22dμ(x)\|f\|_{L^{2}(\mu)}^{2}=\int\|f(x)\|_{2}^{2}\ \mathrm{d}\mu(x)和内积⟨⋅,⋅⟩L2(μ)\langle\cdot,\cdot\rangle_{L^{2}(\mu)}。给定T:Rd→Rd\mathrm{T}:\mathbb{R}^{d}\to\mathbb{R}^{d}，T#μ∈P2(Rd)\mathrm{T}_{\#}\mu\in\mathcal{P}_{2}(\mathbb{R}^{d})是μ\mu的推前测度。

### 2 背景

我们首先回顾Wasserstein空间上优化的一些基本事实。更具体地，我们回顾Wasserstein梯度的概念、Wasserstein空间上的完全凸性，以及(P2(Rd),W2)(\mathcal{P}_{2}(\mathbb{R}^{d}),\mathrm{W}_{2})上的一些经典优化方案。有关Wasserstein梯度流的更多细节，可参见，例如，(Ambrosio等人，2008 (https://arxiv.org/html/2606.27767#bib.bib31)；Lanzetti等人，2025 (https://arxiv.org/html/2606.27767#bib.bib30))。然后，我们简要介绍欧几里得空间上的凸凹过程。

##### Wasserstein梯度。  
令F:P2(Rd)→R\mathcal{F}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\mathbb{R}是一个泛函。它在μ∈P2(Rd)\mu\in\mathcal{P}_{2}(\mathbb{R}^{d})处的Wasserstein梯度∇W2F(μ)∈L2(μ)\nabla_{\mathrm{W}_{2}}\mathcal{F}(\mu)\in L^{2}(\mu)定义为：若对所有ν∈P2(Rd)\nu\in\mathcal{P}_{2}(\mathbb{R}^{d})和γ∈Πo(μ,ν)\gamma\in\Pi_{o}(\mu,\nu)，满足以下一阶泰勒展开（Bonnet，2019 (https://arxiv.org/html/2606.27767#bib.bib32)；Lanzetti等人，2025 (https://arxiv.org/html/2606.27767#bib.bib30)）：

\[
\mathcal{F}(\nu)=\mathcal{F}(\mu)+\int\langle\nabla_{\mathrm{W}_{2}}\mathcal{F}(\mu)(x),y-x\rangle\ \mathrm{d}\gamma(x,y)+o\big(\mathrm{W}_{2}(\mu,\nu)\big). \tag{1}
\]

当梯度存在时，它在L2(μ)L^{2}(\mu)中一般不唯一。然而，只有一个梯度位于切空间TμP2(Rd)⊂L2(μ)T_{\mu}\mathcal{P}_{2}(\mathbb{R}^{d})\subset L^{2}(\mu)中，该空间是Hilbert空间。因此，根据Hilbert分解定理，任何梯度可以分解为TμP2(Rd)T_{\mu}\mathcal{P}_{2}(\mathbb{R}^{d})中的部分和一个正交部分ξ(μ)\xi(\mu)，满足∫⟨ξ(μ),y−x⟩dγ(x,y)=0\int\langle\xi(\mu),y-x\rangle\ \mathrm{d}\gamma(x,y)=0，参见(Lanzetti等人，2025 (https://arxiv.org/html/2606.27767#bib.bib30)，命题2.11)。因此，不失一般性，我们始终使用唯一的∇W2F(μ)∈TμP2(Rd)\nabla_{\mathrm{W}_{2}}\mathcal{F}(\mu)\in T_{\mu}\mathcal{P}_{2}(\mathbb{R}^{d})用于Wasserstein可微泛函，并简称此类泛函为W-可微。

从P2(Rd)\mathcal{P}_{2}(\mathbb{R}^{d})到R\mathbb{R}的经典泛函包括势能V(μ)=∫Vdμ\mathcal{V}(\mu)=\int\mathrm{V}\mathrm{d}\mu和相互作用能W(μ)=12∫∫W(x−y)dμ(x)dμ(y)\mathcal{W}(\mu)=\frac{1}{2}\iint\mathrm{W}(x-y)\ \mathrm{d}\mu(x)\mathrm{d}\mu(y)，其中V,W:Rd→R\mathrm{V},\mathrm{W}:\mathbb{R}^{d}\to\mathbb{R}且W\mathrm{W}对称。若V\mathrm{V}和W\mathrm{W}足够光滑可微，则它们都是可微的（Lanzetti等人，2025 (https://arxiv.org/html/2606.27767#bib.bib30)），其Wasserstein梯度分别为∇W2V(μ)=∇V\nabla_{\mathrm{W}_{2}}\mathcal{V}(\mu)=\nabla\mathrm{V}和卷积∇W2W(μ)=∇W∗μ\nabla_{\mathrm{W}_{2}}\mathcal{W}(\mu)=\nabla\mathrm{W}*\mu。在本工作中，我们主要关注作为相互作用能与势能之和的泛函，因为这些泛函尤其包括了最大均值差异（Arbel等人，2019 (https://arxiv.org/html/2606.27767#bib.bib22)）。

##### Wasserstein空间上的凸性。  
令μ∈P2(Rd)\mu\in\mathcal{P}_{2}(\mathbb{R}^{d})，T,S∈L2(μ)\mathrm{T},\mathrm{S}\in L^{2}(\mu)。给定凸且Gateaux可微的φ:L2(μ)→R\phi:L^{2}(\mu)\to\mathbb{R}，L2(μ)L^{2}(\mu)上T\mathrm{T}和S\mathrm{S}之间的Bregman散度定义为（Frigyik等人，2008 (https://arxiv.org/html/2606.27767#bib.bib33)）：

\[
\mathrm{D}_{\phi}(\mathrm{T},\mathrm{S})=\phi(\mathrm{T})-\phi(\mathrm{S})-\langle\nabla\phi(\mathrm{S}),\mathrm{T}-\mathrm{S}\rangle_{L^{2}(\mu)}.
\]