[논문 리뷰] Fast Convergence for Langevin Diffusion with Manifold Structure

Ankur Moitra, Andrej Risteski|arXiv (Cornell University)|2020. 02. 13.

Markov Chains and Monte Carlo Methods참고 문헌 33인용 수 4

한 줄 요약

이 논문은 행렬 분해 문제에서의 직교 불변성과 같은 대칭성으로 인해 발생하는 다양체 구조를 가진 비로그볼록 분포에서 랭제빈 확산의 빠른 혼합 시간 경계를 수립한다. 수준 집합의 기하학적 구조를 분석하고 확률적 미분 방정식을 활용하여, 미묘한 조건 하에서 다항 시간 샘플링 보장을 입증함으로써, 볼록성의 범위를 초월한 대칭적 비볼록 설정에서의 빠른 혼합을 확장한다. 이는 베이지안 추론과 에너지 기반 모델에서 중요한 응용을 가진다.

ABSTRACT

In this paper, we study the problem of sampling from distributions of the form p(x) \propto e^{-βf(x)} for some function f whose values and gradients we can query. This mode of access to f is natural in the scenarios in which such problems arise, for instance sampling from posteriors in parametric Bayesian models. Classical results show that a natural random walk, Langevin diffusion, mixes rapidly when f is convex. Unfortunately, even in simple examples, the applications listed above will entail working with functions f that are nonconvex -- for which sampling from p may in general require an exponential number of queries. In this paper, we focus on an aspect of nonconvexity relevant for modern machine learning applications: existence of invariances (symmetries) in the function f, as a result of which the distribution p will have manifolds of points with equal probability. First, we give a recipe for proving mixing time bounds for Langevin diffusion as a function of the geometry of these manifolds. Second, we specialize our arguments to classic matrix factorization-like Bayesian inference problems where we get noisy measurements A(XX^T), X \in R^{d imes k} of a low-rank matrix, i.e. f(X) = \|A(XX^T) - b\|^2_2, X \in R^{d imes k}, and βthe inverse of the variance of the noise. Such functions f are invariant under orthogonal transformations, and include problems like matrix factorization, sensing, completion. Beyond sampling, Langevin dynamics is a popular toy model for studying stochastic gradient descent. Along these lines, we believe that our work is an important first step towards understanding how SGD behaves when there is a high degree of symmetry in the space of parameters the produce the same output.

연구 동기 및 목표

• 베이지안 추론과 에너지 기반 모델에서 발생하는 비로그볼록 분포로부터 효율적인 샘플링을 해결하기 위한 도전 과제를 다루는 것.
• 비볼록성에도 불구하고 빠른 혼합을 가능하게 하는 구조적 조건—특히 확률이 동일한 다양체로 인한 대칭성—을 규명하는 것.
• 분포의 수준 집합 기하학과 랭제빈 확산의 혼합 시간을 연결하는 일반적 프레임워크를 개발하는 것.
• 직교 불변성을 가진 행렬 분해 유사 문제에 대해 증명 가능하고 다항 시간 내의 샘플링 보장을 수립하는 것.

제안 방법

• 분포의 수준 집합과 대칭 다양체의 리만 기하학을 바탕으로 랭제빈 확산의 혼합 시간을 경계하는 일반적인 방법을 유도한다.
• 연속 시간 랭제빈 SDE $ dX_t = -\beta \nabla f(X_t) dt + \sqrt{2} dB_t $ 를 분석하여 곡률과 대칭성이 혼합 시간을 감소시킴을 보여준다.
• 이토의 보조정리(Itô’s lemma)를 사용하여 과정의 모멘트를 경계함으로써 $ \mathbb{E}[\|X_t\|_F^p] \leq \text{poly}(d, \|M\|_F) $ 를 증명하며, 이는 이산화 오차를 제어한다.
• 커플링 방법을 통해 이산화 오차 경계를 확립하여 $ \text{KL}(\hat{p}_T, p_T) \leq \text{poly}(d, \|M\|_F) \cdot h^2 $ 를 도출한다. 여기서 $ h $ 는 스텝 크기이다.
• 행렬 분해 문제 $ f(X) = \|\mathcal{A}(XX^T) - b\|_2^2 $ 에 프레임워크를 적용하여 직교 변환에 대해 불변임을 보여준다.
• 핀스커의 부등식(Pinsker’s inequality)과 모멘트 경계를 활용하여, $ f $ 가 다양체 대칭성을 가질 경우 이산화 랭제빈 동역학이 다항 시간 내에 혼합됨을 증명한다.

실험 결과

연구 질문

• RQ1 비로그볼록 분포에서 대칭성에 의해 유도된 동일 확률의 다양체를 가질 경우, 랭제빈 동역학이 빠른 혼합을 달성할 수 있는가?
• RQ2 수준 집합과 대칭 다양체의 기하학이 랭제빈 확산의 혼합 시간에 어떤 영향을 미치는가?
• RQ3 행렬 다양체 구조와 같은 함수 $ f $ 에 대한 어떤 조건이 비볼록성에도 불구하고 증명 가능한 빠른 샘플링을 가능하게 하는가?
• RQ4 스텝 크기와 기저 기하학에 대해 이산화 랭제빈 동역학의 혼합 시간을 어느 정도 경계할 수 있는가?
• RQ5 행렬 분해 문제에서의 직교 불변성이 샘플링 알고리즘의 수렴에 어떤 영향을 미치는가?

주요 결과

• 논문은 대칭 다양체의 리만 기하학과 랭제빈 확산의 혼합 시간을 연결하는 일반적 프레임워크를 수립한다.
• 직교 불변성을 가진 행렬 분해 문제에서는 랭제빈 동역학의 혼합 시간이 차원과 문제 매개변수에 대해 다항식이다.
• 저자들은 $ \mathbb{E}[\|X_t\|_F^p] \leq \|X_0\|_F^p e^{-\frac{4}{5}t \frac{p}{d^2}} + 25d^2\left(\frac{\|M\|_F}{p} + \frac{1}{\beta^2}dp\right) $ 를 증명하여 모멘트 제어를 보장한다.
• 이산화 오차는 $ \text{KL}(\hat{p}_T, p_T) \leq \text{poly}(d, \|M\|_F) \cdot h^2 $ 로 경계되며, 이는 다항 시간 샘플링을 가능하게 한다.
• 이 프레임워크는 행렬 완성, 행렬 감지, 노이즈가 있는 행렬 분해와 같은 구체적 문제에 적용 가능하며, 모두 증명 가능하고 빠른 혼합을 보인다.
• 결과는 매개변수 공간 내의 대칭성과 곡률이 $ f $ 가 비볼록일지라도 효율적인 샘플링을 가능하게 할 수 있음을 시사하며, 비로그볼록성은 비가역성을 암시한다는 전통적 견해에 도전한다.