[Paper Review] Non-Convex Optimization via Non-Reversible Stochastic Gradient Langevin Dynamics

Yuanhan Hu, Xiaoyu Wang+3 more|arXiv (Cornell University)|Apr 6, 2020

Blind Source Separation Techniques34 references4 citations

TL;DR

This paper proposes Non-Reversible Stochastic Gradient Langevin Dynamics (NSGLD), a non-convex optimization method that accelerates convergence by introducing an anti-symmetric matrix to the gradient drift, enabling faster mixing to the stationary distribution. The authors establish non-asymptotic finite-time performance bounds for both population and empirical risk minimization, showing improved convergence rates over standard SGLD with proper matrix selection.

ABSTRACT

Stochastic Gradient Langevin Dynamics (SGLD) is a powerful algorithm for optimizing a non-convex objective, where a controlled and properly scaled Gaussian noise is added to the stochastic gradients to steer the iterates towards a global minimum. SGLD is based on the overdamped Langevin diffusion which is reversible in time. By adding an anti-symmetric matrix to the drift term of the overdamped Langevin diffusion, one gets a non-reversible diffusion that converges to the same stationary distribution with a faster convergence rate. In this paper, we study the non reversible Stochastic Gradient Langevin Dynamics (NSGLD) which is based on discretization of the non-reversible Langevin diffusion. We provide finite-time performance bounds for the global convergence of NSGLD for solving stochastic non-convex optimization problems. Our results lead to non-asymptotic guarantees for both population and empirical risk minimization problems. Numerical experiments for Bayesian independent component analysis and neural network models show that NSGLD can outperform SGLD with proper choices of the anti-symmetric matrix.

Motivation & Objective

• To address the slow convergence of standard Stochastic Gradient Langevin Dynamics (SGLD) in non-convex optimization.
• To develop a non-reversible variant of SGLD that preserves the same stationary distribution but converges faster.
• To provide finite-time, non-asymptotic performance guarantees for global convergence in both population and empirical risk minimization.
• To analyze the impact of the anti-symmetric matrix on convergence speed and stability.
• To validate the method empirically on Bayesian ICA and neural network models.

Proposed method

• NSGLD is derived from the Euler discretization of a non-reversible Langevin SDE with drift term modified by an anti-symmetric matrix $ J $, forming $ A_J = I + J $.
• The algorithm updates parameters via $ X_{k+1} = X_k - \eta A_J g_k + \sqrt{2\eta \beta^{-1}} \xi_k $, where $ g_k $ is a stochastic gradient estimator.
• The stationary distribution of the non-reversible SDE matches that of the original overdamped Langevin SDE, ensuring correct sampling of the target distribution.
• Finite-time performance bounds are derived using spectral gap analysis and uniform stability, incorporating the effect of $ J $ on convergence rate.
• The method leverages non-reversible dynamics to improve the spectral gap, leading to faster mixing and reduced convergence time.
• Theoretical bounds are established for both population and empirical risk minimization, with explicit dependence on $ \beta $, $ d $, $ J $, and the spectral gap $ \lambda_{\mathbf{z},J} $.

Experimental results

Research questions

• RQ1 Can non-reversible dynamics improve the convergence rate of stochastic gradient Langevin dynamics in non-convex optimization?
• RQ2 What is the finite-time performance guarantee of NSGLD for global convergence in population and empirical risk minimization?
• RQ3 How does the choice of the anti-symmetric matrix $ J $ affect the convergence speed and stability of NSGLD?
• RQ4 What is the explicit dependence of the convergence bound on key parameters such as $ \beta $, $ d $, and the spectral gap?
• RQ5 Does NSGLD outperform SGLD in practical non-convex optimization tasks like Bayesian ICA and neural network training?

Key findings

• NSGLD achieves faster convergence than SGLD by leveraging non-reversible dynamics, which increases the spectral gap of the underlying diffusion process.
• The finite-time performance bound for NSGLD is $ \tilde{\mathcal{O}}\left(\hat{C}_{\mathbf{z},J}\varepsilon + \frac{\sqrt{\beta}(\beta+d)}{\sqrt{\lambda_{\mathbf{z},J=0}}}\varepsilon + \frac{d\log(1+\beta)}{\beta}\right) $, with improved dependence on $ \lambda_{\mathbf{z},J=0} $.
• The spectral gap $ \lambda_{\mathbf{z},J} $ for the non-reversible SDE satisfies $ \lambda_{\mathbf{z},J} > \lambda_{\mathbf{z},J=0} $, ensuring strictly faster convergence than the reversible case.
• Numerical experiments show NSGLD outperforms SGLD in Bayesian independent component analysis and neural network training with properly chosen $ J $.
• The uniform stability bound between two datasets differing in one sample is $ \tilde{\mathcal{O}}\left(\frac{1}{n}\right) $, with explicit dependence on $ \beta $, $ d $, and $ \lambda_* $.
• The norm of the matrix $ A_J $ satisfies $ \|A_J\|^2 = 1 + \|J\|^2 \geq 1 $, which influences the gradient scaling and convergence behavior.