[论文解读] Fast Convergence for Langevin Diffusion with Manifold Structure

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

Markov Chains and Monte Carlo Methods参考文献 33被引用 4

一句话总结

本文建立了在具有对称性结构（如矩阵分解问题中的正交不变性）的非对数凸分布中，Langevin扩散过程的快速混合时间界限。通过分析等概率集的几何结构并利用随机微分方程，作者在温和条件下证明了多项式时间采样保证，将快速混合的适用范围从凸性扩展到与贝叶斯推断和能量模型相关的对称非凸设置。

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.

研究动机与目标

• 解决在贝叶斯推断和能量模型中出现的非对数凸分布的高效采样挑战。
• 识别结构性条件——特别是由对称性诱导的等概率流形——以实现尽管存在非凸性仍能快速混合。
• 建立一个通用框架，将分布等高线集的几何结构与Langevin扩散的混合时间联系起来。
• 为具有正交不变性的矩阵分解类问题建立可证明的多项式时间采样保证。

提出的方法

• 基于分布等高线集的黎曼几何与对称性流形，推导出Langevin扩散混合时间的通用界 bound 方法。
• 分析连续时间Langevin SDE $ dX_t = -\beta \nabla f(X_t) dt + \sqrt{2} dB_t $，表明曲率与对称性可降低混合时间。
• 使用伊藤引理控制过程的矩，证明 $ \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不等式与矩界，证明当 $ f $ 具有流形对称性时，离散化Langevin动力学可在多项式时间内混合。

实验结果

研究问题

• RQ1 Langevin动力学是否能在表现出对称性诱导的等概率流形的非对数凸分布中实现快速混合？
• RQ2 等高线集与对称性流形的几何结构如何影响Langevin扩散的混合时间？
• RQ3 函数 $ f $ 的哪些条件（如矩阵流形结构）可允许在非凸性下实现可证明的快速采样？
• RQ4 离散Langevin动力学的混合时间在多大程度上可基于步长与底层几何结构进行界定？
• RQ5 矩阵分解问题中正交不变性的存在如何影响采样算法的收敛性？

主要发现

• 本文建立了一个通用框架，将对称性流形的黎曼几何与Langevin扩散的混合时间联系起来。
• 对于具有正交不变性的矩阵分解问题，Langevin动力学的混合时间在维度与问题参数上为多项式时间。
• 作者证明了 $ \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 $ 非凸时仍实现高效采样，挑战了传统观点中非对数凸性意味着计算不可行的观念。