Constrained Exploration via Reflected Replica Exchange Stochastic Gradient Langevin Dynamics

TL;DR

This paper introduces reflected reSGLD, a constrained sampling algorithm that improves exploration efficiency in non-convex spaces by incorporating reflection steps within bounded domains, supported by theoretical and empirical evidence.

Contribution

We propose reflected reSGLD, a novel constrained sampling method that enhances mixing rates through reflection steps, addressing stagnation issues in high-temperature chains.

Findings

Reducing domain diameter quadratically improves mixing rates.

Reflected reSGLD effectively explores constrained non-convex distributions.

Empirical results demonstrate improved performance in physical systems and image classification.

Abstract

Replica exchange stochastic gradient Langevin dynamics (reSGLD) is an effective sampler for non-convex learning in large-scale datasets. However, the simulation may encounter stagnation issues when the high-temperature chain delves too deeply into the distribution tails. To tackle this issue, we propose reflected reSGLD (r2SGLD): an algorithm tailored for constrained non-convex exploration by utilizing reflection steps within a bounded domain. Theoretically, we observe that reducing the diameter of the domain enhances mixing rates, exhibiting a $\textit{quadratic}$ behavior. Empirically, we test its performance through extensive experiments, including identifying dynamical systems with physical constraints, simulations of constrained multi-modal distributions, and image classification tasks. The theoretical and empirical findings highlight the crucial role of constrained exploration in improving the simulation efficiency.