Replica Exchange for Non-Convex Optimization

Jing Dong and Xin T. Tong † January 24, 2020 Abstract Gradient descent (GD) is known to converge quickly for convex objective functions, but it can be trapped at local minimums. On the other hand, Langevin dynamics (LD) can explore the state space and find global minimums, but in order to give accurate estimates, LD needs to run with small discretization stepsize and weak stochastic force, which in general slow down its convergence. This paper shows that these two algorithms can "collaborate" through a simple exchange mechanism, in which they swap their current positions if LD yields a lower objective function. This idea can be seen as the singular limit of the replica exchange technique from the sampling literature. We show that this new algorithm converges to the global minimum linearly with high probability, assuming the objective function is strongly convex in a neighborhood of the unique global minimum. By replacing gradients with stochastic gradients, and adding a proper threshold to the exchange mechanism, our algorithm can also be used in online settings. We further verify our theoretical results through some numerical experiments, and observe superior performance of the proposed algorithm over running GD or LD alone. 1 Introduction Division of labor is the secret of any efficient enterprises. By collaborating with individuals with different skillsets, we can focus on tasks within our own expertise and produce better outcomes than working independently. This paper asks whether the same principle can be applied when designing an algorithm. Given a general smooth non-convex objective function F, we consider the unconstrained optimization problem min x R dF ( x). However, this local minimum may not be the global minimum, and GD will be trapped there afterwards. On the other hand, sampling-based algorithms, such as the Langevin dynamics (LD) can escape local minimums by their stochasticity, but the additional stochastic noise contaminates the optimization results and slows down the convergence when the iterate is near the global minimum. In general, deterministic algorithms are designed to finding local minimums quickly, but they can be terrible in exploration. Sampling-based algorithms are better suited for exploring the state space, but they are inefficient when pinpointing the local minimums. This paper investigates how they can "collaborate" to get the "best of the two worlds". The collaboration mechanism we introduced here comes from replica-exchange in the sampling literature. Its implementation is very simple: we run a copy of GD, denoted by X n; and a copy of LD, denoted by Y n. If F (X n) F (Y n), we swap their positions.