[논문 리뷰] Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

Jason M. Altschuler, Sinho Chewi외 2명|arXiv (Cornell University)|2021. 06. 15.

Geometric Analysis and Curvature Flows인용 수 8

한 줄 요약

이 논문은 워셔스타인 거리의 바리센터, 엔트로피 정규화된 바리센터, 그리고 가우시안 분포의 기하학적 중앙값을 계산할 때 리만 기울기 하강법(Riemannian gradient descent)이 차원에 영향을 받지 않는 수렴 속도를 확립한다. 저자들은 새로운 지오데식 볼록성 결과를 증명함으로써, 차원에 의존하지 않는 수렴 보장을 달성하여 최적 운반 비용 거리에 대한 최적화 이론적 격차를 해결한다.

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.

연구 동기 및 목표

• 최적 운반 비용 바리센터에 대한 리만 기울기 하강법에서 기존 수렴 속도가 차원에 따라 지수적으로 증가하는 이론적 격차를 해소한다.
• 가우시안 분포의 워셔스타인 바리센터 맥락에서 리만 기울기 하강법에 대해 차원에 영향을 받지 않는 수렴 보장을 제공한다.
• 기존에 리만 기울기 하강법에 대한 수렴 보장이 없었던 두 가지 관련 문제인 엔트로피 정규화된 바리센터와 기하학적 중앙값 문제로 분석을 확장하여, 이들 문제에 대해 처음으로 수렴 결과를 제시한다.
• 더 강력한 지오데식 볼록성 제어를 통해, 리만 기하학적 설정에서 목적 함수의 비볼록성에 기인한 문제를 극복하고 차원에 영향을 받지 않는 수렴을 달성한다.
• 실제로 리만 기울기 하강법이 이론적 한계가 있는 타 방법들보다 우수한 성능을 보임을 보이며, 이는 이론적으로도 설명 가능하다.

제안 방법

• 워셔스타인 바리센터 문제를 양의 정부호 행렬의 콘과 등장하는 Bures–Wasserstein 다양체 상의 최소화 문제로 재구성한다.
• 수렴을 제어하고 차원에 영향을 받지 않는 수렴을 보장하기 위해 목적 함수의 새로운 지오데식 볼록성 성질을 도입하고 증명한다.
• 공분산 행렬의 다양체 상에서 직접 리만 기울기 하강법을 적용하여, 유클리드 공간에 통합하는 대신 다양체의 내재 기하학을 활용한다.
• 엔트로피 정규화된 바리센터와 기하학적 중앙값 문제를 동일한 기하학적 프레임워크를 사용하여 분석함으로써, 표준 바리센터를 넘어서 수렴 보장을 확장한다.
• 지오데식 지도와 항등원에서의 지수 지도 및 로그 지도를 포함한 최적 운반 이론 및 리만 기하학 도구를 사용하여 수렴 행동을 분석한다.
• 환경 차원 d에 의존하지 않는 이론적 수렴 속도를 제공하며, 다양체 상의 곡률과 볼록성 제어에 기반한다.

실험 결과

연구 질문

• RQ1 목적 함수가 지오데식 비볼록일 경우에도 리만 기울기 하강법이 가우시안 분포의 워셔스타인 바리센터에 대해 차원에 영향을 받지 않는 수렴을 달성할 수 있는가?
• RQ2 Bures–Wasserstein 다양체의 기하학적 성질과 볼록성 특성은 고차원 환경에서 차원에 영향을 받지 않는 수렴을 가능하게 하는가?
• RQ3 기존에 리만 기울기 하강법에 대한 수렴 보장이 없었던 엔트로피 정규화된 바리센터와 기하학적 중앙값 문제로 동일한 이론적 프레임워크를 확장할 수 있는가?
• RQ4 리만 기울기 하강법이 실무에서 유럽 기울기 하강법과 SDP 솔버보다 우수한 성능을 보이는 이유는 무엇이며, 이는 이론적으로 설명될 수 있는가?
• RQ5 다양체의 곡률과 지오데식 구조는 수렴의 안정성과 가속화에 어떤 역할을 하는가?

주요 결과

• 이 논문은 가우시안 분포의 워셔스타인 바리센터에 대해 Bures–Wasserstein 다양체 상에서 리만 기울기 하강법에 대해 처음으로 차원에 영향을 받지 않는 수렴 속도를 확립한다.
• 수렴을 더 강력하게 제어할 수 있도록 새로운 지오데식 볼록성 결과를 증명하여, 수렴이 차원 d에 영향을 받지 않도록 보장한다.
• 엔트로피 정규화된 바리센터와 기하학적 중앙값 문제로 수렴 보장을 확장하여, 이들 문제에 대해 리만 기울기 하강법에 대해 처음으로 이러한 결과를 제시한다.
• 실증 결과로 리만 기울기 하강법이 스텝 사이즈 조정 없이도 유클리드 기울기 하강법과 SDP 솔버보다 더 빠르게 수렴하는 것으로 나타났으며, 후자의 이론적 한계는 차원에 따라 지수적으로 증가한다.
• 이론적 분석을 통해 리만 기울기 하강법이 악조건의 행렬과 행렬의 구조(예: 비가환 고유기저)에 대해 안정적으로 수렴함을 확인하였으며, 다양한 합성 데이터 분포에서 안정적인 수렴을 보였다.
• 실험에서는 고정된 스텝 사이즈 1로도 수렴을 달성하였고, 유클리드 기울기 하강법는 신중한 조정이 필요하나 여전히 성능이 열 劣하다. 이는 내재 기하학의 우수성을 강조한다.