[论文解读] Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

Jason M. Altschuler, Sinho Chewi+2 more|arXiv (Cornell University)|Jun 15, 2021

Geometric Analysis and Curvature Flows被引用 8

一句话总结

该论文在计算高斯分布的Wasserstein均值、熵正则化均值和几何中位数时，建立了黎曼梯度下降在Bures–Wasserstein流形上的无维收敛速率。通过证明新的测地线凸性结果，作者实现了与维度无关的收敛保证，解决了最优传输度量下优化领域中的一个关键理论空白。

ABSTRACT

We study first-order optimization algorithms for computing the barycenter of Gaussian distributions with respect to the optimal transport metric. Although the objective is geodesically non-convex, Riemannian GD empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean GD and SDP solvers. This stands in stark contrast to the best-known theoretical results for Riemannian GD, which depend exponentially on the dimension. In this work, we prove new geodesic convexity results which provide stronger control of the iterates, yielding a dimension-free convergence rate. Our techniques also enable the analysis of two related notions of averaging, the entropically-regularized barycenter and the geometric median, providing the first convergence guarantees for Riemannian GD for these problems.

研究动机与目标

• 解决黎曼梯度下降在最优传输均值问题中的理论空白，其中现有收敛速率对维度呈指数依赖。
• 为高斯分布Wasserstein均值的黎曼梯度下降提供无维收敛保证。
• 将分析扩展至两个相关问题：熵正则化均值和几何中位数，首次为这些问题提供黎曼梯度下降的收敛结果。
• 建立更强的测地线凸性控制，以实现无维收敛，克服黎曼设置下目标函数的非凸性。
• 表明尽管存在理论局限性，黎曼梯度下降在实践中优于欧氏梯度下降和SDP求解器。

提出的方法

• 将Wasserstein均值问题表述为在Bures–Wasserstein流形上的最小化问题，该流形与正半定矩阵的锥等距。
• 引入并证明目标函数的新测地线凸性性质，以控制迭代点并确保无维收敛。
• 直接在协方差矩阵的流形上应用黎曼梯度下降，利用流形的内在几何结构，而非嵌入到欧氏空间中。
• 使用相同的几何框架分析熵正则化均值和几何中位数问题，将收敛保证扩展至标准均值之外。
• 利用最优传输和黎曼几何的工具，包括单位元处的指数映射和对数映射，分析收敛行为。
• 提供与环境维度d无关的理论收敛速率，依赖于流形上的曲率和凸性控制。

实验结果

研究问题

• RQ1 尽管目标函数在测地线上非凸，黎曼梯度下降是否能为高斯分布的Wasserstein均值实现无维收敛？
• RQ2 Bures–Wasserstein流形的哪些几何和凸性性质使得在高维设置下可实现无维收敛？
• RQ3 该理论框架是否可扩展至熵正则化均值和几何中位数问题，而这些问题此前缺乏黎曼梯度下降的收敛保证？
• RQ4 为何黎曼梯度下降在实践中优于欧氏梯度下降和SDP求解器？这一现象能否从理论上得到解释？
• RQ5 流形的曲率和测地线结构在稳定和加速收敛过程中起到何种作用？

主要发现

• 该论文首次在Bures–Wasserstein流形上为高斯分布的Wasserstein均值建立了无维收敛速率。
• 证明了新的测地线凸性结果，为迭代点提供了更强的控制，实现了与维度d无关的收敛。
• 收敛保证被扩展至熵正则化均值和几何中位数，首次为黎曼梯度下降在这两个问题上提供了此类结果。
• 实验结果表明，黎曼梯度下降在无需步长调优的情况下，收敛速度仍快于欧氏梯度下降和SDP求解器，而后者理论边界的依赖维度呈指数增长。
• 理论分析证实，黎曼梯度下降对病态条件和矩阵结构（如非对易特征基）具有鲁棒性，在多种合成数据分布下均表现出稳定收敛。
• 实验中采用固定步长1即可实现收敛，而欧氏梯度下降需精细调优且性能仍较差，凸显了内在几何结构的优势。