[论文解读] Gradient descent algorithms for Bures-Wasserstein barycenters

Sinho Chewi, Tyler Maunu+2 more|arXiv (Cornell University)|Jan 6, 2020

Geometric Analysis and Curvature Flows参考文献 33被引用 28

一句话总结

本文为计算概率测度的Bures-Wasserstein中位数开发了梯度下降算法，尽管缺乏测地线凸性，仍通过利用Polyak-Łojasiewicz（PL）不等式建立了全局收敛速率。主要贡献在于证明了高斯测度的Bures-Wasserstein流形上的PL不等式，从而首次在该设定下实现了对一阶方法的全局收敛速率分析。

ABSTRACT

We study first order methods to compute the barycenter of a probability distribution $P$ over the space of probability measures with finite second moment. We develop a framework to derive global rates of convergence for both gradient descent and stochastic gradient descent despite the fact that the barycenter functional is not geodesically convex. Our analysis overcomes this technical hurdle by employing a Polyak-Lojasiewicz (PL) inequality and relies on tools from optimal transport and metric geometry. In turn, we establish a PL inequality when $P$ is supported on the Bures-Wasserstein manifold of Gaussian probability measures. It leads to the first global rates of convergence for first order methods in this context.

研究动机与目标

• 开发用于在Bures-Wasserstein几何中计算Wasserstein中位数的一阶优化方法。
• 解决在度量空间中非凸设定下梯度下降缺乏理论收敛保证的问题。
• 在中位数泛函缺乏测地线凸性的情况下，仍建立全局收敛速率。
• 证明高斯测度的Bures-Wasserstein流形上的Polyak-Łojasiewicz（PL）不等式。
• 在这一非欧几里得优化背景下，首次实现对随机与标准梯度下降的全局收敛速率分析。

提出的方法

• 推导出在高斯测度的Bures-Wasserstein流形上，中位数泛函的Polyak-Łojasiewicz（PL）不等式。
• 利用最优传输和度量几何的工具，分析非凸设定下的收敛性。
• 应用PL不等式，建立梯度下降与随机梯度下降的全局线性收敛速率。
• 采用广义测地线与Monge-Ampère方程，分析测度空间中路径上密度的行为。
• 以传输映射与协方差矩阵的形式刻画中位数泛函的Wasserstein梯度。
• 证明该泛函在测度的对数密度沿广义测地线方向具有凸性，从而支持基于曲率的分析。

实验结果

研究问题

• RQ1 尽管中位数泛函缺乏测地线凸性，是否仍可为Bures-Wasserstein中位数泛函的梯度下降建立全局收敛速率？
• RQ2 中位数泛函在高斯测度的Bures-Wasserstein流形上是否满足Polyak-Łojasiewicz（PL）不等式？
• RQ3 在该设定下，计算Wasserstein中位数的梯度下降与随机梯度下降的全局收敛速率是多少？
• RQ4 Bures-Wasserstein流形的几何结构如何促成非凸优化中收敛保证的推导？
• RQ5 PL不等式能否用于弥合中位数计算中统计一致性与算法效率之间的差距？

主要发现

• 中位数泛函在$ \zeta $-正则的高斯测度Bures-Wasserstein流形上满足PL不等式，常数为$ C_{\mathsf{PL}} = \zeta^2/4 $。
• 在Bures-Wasserstein中位数问题上，梯度下降与随机梯度下降均建立了全局线性收敛速率。
• PL不等式对所有$ b \in \mathcal{S}_\zeta $一致成立，其中$ \mathcal{S}_\zeta $为特征值位于$[\zeta, 1]$内的中心高斯测度集合。
• 收敛速率依赖于最小特征值$ \zeta $，条件越好收敛越快。
• 分析表明，泛函$ \rho \mapsto \ln\|\rho\|_{L^\infty} $沿广义测地线是凸的，这是关键的技术结果。
• 本研究首次在高斯测度全流形的背景下，为一阶方法提供了Bures-Wasserstein中位数的全局收敛保证。