Title: ACCELERATING NONCONVEX LEARNING VIA REPLICA EXCHANGE LANGEVIN DIFFUSION

Abstract: Langevin diffusion is a powerful method for nonconvex optimization, which enables the escape from local minima by injecting noise into the gradient. In particular, the temperature parameter controlling the noise level gives rise to a tradeoff between "global exploration" and "local exploitation", which correspond to high and low temperatures. To attain the advantages of both regimes, we propose to use replica exchange, which swaps between two Langevin diffusions with different temperatures. We theoretically analyze the acceleration effect of replica exchange from two perspectives: (i) the convergence in χ 2 -divergence, and (ii) the large deviation principle. Such an acceleration effect allows us to faster approach the global minima. Furthermore, by discretizing the replica exchange Langevin diffusion, we obtain a discrete-time algorithm. For such an algorithm, we quantify its discretization error in theory and demonstrate its acceleration effect in practice.

1. INTRODUCTION:

We consider the problem of minimizing a nonconvex objective function, which arises from numerous machine learning problems such as training neural networks. However, due to the existence of spurious local minima, nonconvex optimization remains challenging in both theory and practice. To overcome this difficulty, one idea is to construct a diffusion process whose invariant distribution concentrates around the global minima (Chiang et al., 1987) . If we run such a diffusion process for a sufficiently long time, we expect to draw from its invariant distribution a point that is close to a global minimum with high probability.

One commonly used diffusion process is the Langevin diffusion, which is closely related to firstorder optimization (Gidas, 1985) . In specific, the Langevin diffusion can be viewed as gradient flow with injected noise, where the noise is scaled by a temperature parameter. Such a temperature parameter gives rise to a tradeoff between two opposite effects, namely "global exploration" and "local exploitation". More specifically, a higher temperature results in a larger injected noise, which increases the probability of escaping from spurious local minima and facilitates the global exploration of the whole domain. On the other hand, with a lower temperature, the Langevin diffusion behaves more like the gradient flow, which focuses more on exploiting the local geometry to decrease the objective value, yielding a more concentrated invariant distribution.

In this paper, we aim to bridge the gap between "global exploration" and "local exploitation" in nonconvex optimization with replica exchange, which originates from parallel tempering Markov chain Monte Carlo methods for sampling from a multimodal distribution (Earl & Deem, 2005) . In contrast to the standard Langevin diffusion driven by a single temperature, the replica exchange Langevin diffusion consists of two Langevin diffusions with low and high temperatures, which locally exploits the geometry and globally explores the domain, respectively. In particular, the two Langevin diffusions exchange their positions following a specific swapping scheme, which ensures that the invariant distribution concentrates around the global minima. See Figures 1.(a)-(b) for an illustration of swapping. Compared with the standard Langevin diffusion, replica exchange accelerates the rate of convergence to the invariant distribution. As a result, with replica exchange we obtain a more accurate solution to the nonconvex optimization problem within a shorter time. It is worth mentioning that, although in this paper we focuses on the case of swapping two Langevin diffusions, our theory and method naturally extend to multiple Langevin diffusions.

Our theory quantifies the acceleration effect of replica exchange from two perspectives. First, we use the theory of Markov semigroup, particularly the Dirichlet form, to characterize the evolution of the χ 2 -divergence between the distribution at time t and the invariant distribution. In specific, we show that, compared with the standard Langevin diffusion, replica exchange boosts the Dirichlet form with a nonnegative term, which results in a faster decay of the χ 2 -divergence. Further combined with the Poincaré inequality, we establish the exponential rate of convergence of the χ 2 -divergence to zero, where the acceleration effect of replica exchange is characterized by the Poincaré constant. We illustrate such an acceleration effect in Figure 1 . (c) . Second, we use the large deviation principle (LDP) to characterize the exponential decay of the probability that the empirical measure of the sample path deviates from the invariant distribution. In particular, we show that replica exchange boosts the LDP rate function, which implies a faster exponential decay rate. See Figure 1 .(d) for an illustration. Both perspectives are built on the infinitesimal generator of the replica exchange Langevin diffusion. We show that essentially the acceleration effects are the consequences of an additional term in the infinitesimal generator, which is introduced by replica exchange. In addition to quantifying the acceleration effect of replica exchange, we use the Euler scheme to discretize the replica exchange Langevin diffusion and propose a discrete-time algorithm. We establish an upper bound of the discretization error via Grönwall's inequality. We also discuss several issues related to the parameters of the replica exchange Langevin diffusion, especially the swapping intensity and temperatures. We defer these discussions to §A of the appendix.

Related Work::

The idea of nonconvex optimization via diffusion process dates back to simulated annealing ( Černỳ, 1985) and is systematically studied in Chiang et al. (1987) . For the application of the Langevin diffusion in nonconvex optimization, see, for example, Gidas (1985) and Geman & Hwang (1986) and the followup works. More recently, Dalalyan & Tsybakov (2009) ; Bubeck et al. (2015) ; Dalalyan (2017) analyze the nonasymptotic rate of convergence of the Langevin diffusion for strongly convex objective functions. Moreover, Durmus et al. (2017) ; Zhang et al. (2017) ; Raginsky et al. (2017) ; Xu et al. (2018) further extend the analysis to handle nonconvexity and stochastic gradient. See also, for example, Belloni et al. (2015) ; Ge et al. (2015) ; Hazan et al. (2016) ; Cheng et al. (2017) ; Zou et al. (2018b; a; 2019) for related works on the extensions of the Langevin diffusion and other diffusions. However, all these works focus on the setting with a single temperature. In comparison, we study the acceleration effect of replica exchange, which utilizes more than one temperature. Hence, our work is orthogonal to these existing works and can potentially be applied to accelerate other diffusion processes. Woodard et al. (2009) study the mixing time of parallel tempering, while Dupuis et al. (2012) extend the swapping rate of replica exchange to infinity and establish the corresponding LDP. In addition, another well-studied method is simulated tempering, a technique similar to replica exchange, which also aims at accelerating the convergence of diffusion processes (Marinari & Parisi, 1992; Geyer & Thompson, 1995) . The main difference between replica exchange and simulated tempering is that the former tracks multiple diffusion processes driven by various temperatures and accelerates the convergence by swapping, while the latter tracks only one diffusion process but treats the temperature as a stochastic process. See, for example, Zheng (2003) and the references therein for a detailed discussion on the difference. More recently, Ge et al. (2017) study the convergence of simulating tempering for nonconvex objective functions from the perspective of mixing time, which is also orthogonal to our work. The LDP techniques in our work may potentially be applied to the analysis of simulated tempering as well.

Main Contribution: Despite the broad application in sampling, to the best of our knowledge, this paper is the first attempt to apply the replica exchange Langevin diffusion in the context of nonconvex optimization. In summary, our contribution is twofold: (i) We quantify the acceleration effect of the replica exchange Langevin diffusion from two perspectives, that is, the convergence of the χ 2 -divergence and the LDP for empirical measures. (ii) We propose a nonconvex optimization algorithm based on the discretization of the replica exchange Langevin diffusion and establish an upper bound for the discretization error.

2. BASIC IDEA:

We consider the following unconstrained optimization problem,

minimize x∈R d U (x).

(2.1)

When U (x) is nonconvex, it is difficult to obtain its global minima. A commonly used algorithm is the Langevin diffusion, which is defined by the stochastic differential equation

dX t = -∇U (X t )dt + √ 2τ dW t . (2.2)

Here {W t } t≥0 is a standard d-dimensional Brownian motion and τ > 0 is the temperature parameter. Under mild regularity conditions, the Langevin diffusion {X t } t≥0 has a unique invariant distribution that is absolutely continuous with respect to the Lebesgue measure with density

π τ (x) = e -U (x)/τ ∫ R d e -U (x)/τ dx .

(2.3)

Here ∫ R d e -U (x)/τ dx is the normalization constant. In particular, π τ is the limiting distribution of {X t } t≥0 , which means that with any initialization X 0 , the distribution of X t converges to π τ (Bakry et al., 2013) . We observe from (2.3) that with τ → 0 the probability measure π τ concentrates around the global minima of U (x) (Hwang, 1980) . In other words, if we choose a sufficiently small temperature parameter τ and run the Langevin diffusion in (2.2) for a sufficiently long time, we expect to obtain a solution that falls into the neighborhood of a global minimum of U (x) with high probability.

Here the temperature parameter τ plays a crucial role. In practice, we can only run the Langevin diffusion for a finite time, giving rise to a tradeoff between "global exploration" and "local exploitation", which correspond to high and low temperatures. In specific, when the temperature parameter τ is small, the convergence of X t is slow and the particle can be trapped by a local minimum for a long time without globally exploring the whole domain. Consequently, within a finite time, we can only obtain a sample from a distribution that is far away from the invariant distribution π τ . In contrast, when the temperature parameter τ is large, the convergence is accelerated due to better global exploration. The distribution of X t converges faster to the invariant distribution π τ globally, but locally π τ is less concentrated around the global minima of U (x).

To bridge the gap between high and low temperatures, in this paper we study an adaptive algorithm called replica exchange Langevin diffusion. In detail, we consider a pair of particles driven by two Langevin diffusions as defined in (2.2) with temperatures τ 1 < τ 2 , respectively. We use

Z t = (Z (1) t , Z

(2) t ) to denote the positions of the two particles at time t. In other words, we have

dZ (1) t = -∇U ( Z (1) t ) dt + √ 2τ 1 dW (1) t , dZ (2) t = -∇U ( Z (2) t ) dt + √ 2τ 2 dW (2) t .

(2.4)

According to (2.3), the invariant distribution of {Z t } t≥0 is absolutely continuous with respect to the Lebesgue measure on R 2d and its density is proportional to

µ(x 1 , x 2 ) = exp ( -U (x 1 )/τ 1 -U (x 2 )/τ 2 ) .

(2.5)

The marginal invariant distribution of the low-temperature particle concentrates around the global minima of U (x). Hence, the low-temperature particle is of particular interest for the purpose of nonconvex optimization. The key idea of replica exchange is to enable the low-temperature particle to achieve better global exploration by swapping its position with the high-temperature particle. In specific, at time t and positions (x 1 , x 2 ), the two particles swap with rate

a • s(x 1 , x 2 ) = a • ( 1 ∧ µ(x 2 , x 1 ) µ(x 1 , x 2 ) ) , (2.6)

which means

P ( Z t+dt = (x 2 , x 1 ) | Z t = (x 1 , x 2 ) ) = a • s(x 1 , x 2 )dt, P ( Z t+dt = (x 1 , x 2 ) | Z t = (x 1 , x 2 ) ) = 1 -a • s(x 1 , x 2 )dt. (2.7)

Here a ≥ 0 is a constant called swapping intensity. As is shown in Lemma 3.2, the specific form of s(x 1 , x 2 ) in (2.6) ensures that for any a, the invariant distribution of the replica exchange Langevin diffusion {Z t } t≥0 is the same as (2.5), that is, as if the two particles are independent. Corresponding to the continuous-time process in (2.4)-(2.6), in §3.4 of the appendix we consider the replica exchange stochastic gradient descent algorithm, which corresponds to the discretization of {Z t } t≥0 using Euler scheme.

To better understand the intuition behind replica exchange, note that we use two particles driven by low and high temperatures to achieve "local exploitation" and "global exploration", respectively.

For the purpose of optimization, we only need to track the trajectory of the first particle. By plugging (2.5) in (2.6), we obtain

a • s(x 1 , x 2 ) = a • exp ( 0 ∧ (1/τ 1 -1/τ 2 ) • ( U (x 1 ) -U (x 2 ) ) )

.

(2.8)

Since τ 1 < τ 2 , the rate a • s(x 1 , x 2 ) is monotone increasing with respect to the difference between the objective values at x 1 and x 2 , which denote the positions of the two particles. Hence, the two particles are more likely to swap when the first one has a larger objective value. In other words, swapping tends to move the first particle, which we are interested in, to a position corresponding to a lower objective value. An extreme case of the replica exchange Langevin diffusion is τ 1 = 0 and τ 2 = ∞. In this case, the first equation in (2.4) reduces to an ordinary differential equation charactering the deterministic gradient descent, while the second one in (2.4) corresponds to an approximated uniform exploration of the whole domain R d . According to (2.8), the particles never swap if the objective value of the first particle is smaller than the second. Hence, roughly speaking, the replica exchange Langevin diffusion reduces to the deterministic gradient descent with uniformly randomized restarts. In contrast to this extreme case, for 0 < τ 1 < τ 2 < ∞, the second particle globally explores the whole domain more adaptively with not only noise but also gradient information. In particular, its stationary distribution has a larger density around the local minima, which leads to better restarts that adapt to the global geometry for the first particle.

3. THEORETICAL ANALYSIS:

In this section, we lay out the theoretical analysis of replica exchange Langevin diffusion introduced in §2, which demonstrates the acceleration effect of swapping. Due to space constraint, we defer the necessary background and detailed proofs to §B and §C of the appendix, respectively. Throughout the following analysis, we assume that the objective function U (•) satisfies the following assumption.

Assumption 3.1 (Smoothness and Dissipativity). The function

U (•) is L-smooth, that is, there exists a positive constant L such that for all x, y ∈ R d ∇U (x) -∇U (y) ≤ L∥x -y∥. (3.1)

The function U (•) is also (α, β)-dissipative, that is, there exist positive constants α and β such that for all

x ∈ R d ⟨ x, ∇U (x) ⟩ ≥ α∥x∥ 2 -β. (3.2)

The assumption on smoothness in (3.1) characterizes the Lipschitz continuity of the gradient of objective function, which is commonly used in the optimization literature (Nesterov, 2013) . The assumption on dissipativity in (3.2), roughly speaking, characterizes the approximate quadratic growth of the objective function at infinity, which is commonly used in the control and dynamic system literature (Hale, 2010) . This condition is also used in Raginsky et al. (2017) . It is worth noting that our theory does not require the convexity assumption.

3.1. INVARIANCE AND REVERSIBILITY:

In this section, we show that the invariant distribution of the replica exchange Langevin diffusion {Z t } t≥0 has the form in (2.5) rigorously, whose first component preserves the concentration. We also show the reversibility of {Z t } t≥0 , which is a nice property necessary for the subsequent convergence analysis. We consider the replica exchange Langevin diffusion {Z t } t≥0 defined by (2.4)-(2.6), where

Z t = (Z (1)

t , Z

(2) t ) denotes the positions of the two particles at time t, and a ≥ 0 is the swapping intensity that controls the frequency of swapping. The following auxiliary lemma characterizes the invariant distribution of {Z t } t≥0 and its reversibility. Restricted to the space, we provide a detailed proof in §C.1 of the appendix.

Lemma 3.2. The infinitesimal generator of {Z t } t≥0 with swapping intensity a, which is defined in (2.4)-(2.6), takes the form

L a ( f (x 1 , x 2 ) ) = - ⟨ ∇ x1 f (x 1 , x 2 ), ∇ x1 U (x 1 ) ⟩ + τ 1 ∆ x1 f (x 1 , x 2 ) L a 1 ( f (x 1 , x 2 ) ) (3.3) - ⟨ ∇ x2 f (x 1 , x 2 ), ∇ x2 U (x 2 ) ⟩ + τ 2 ∆ x2 f (x 1 , x 2 ) L a 2 ( f (x 1 , x 2 ) ) + a • s(x 1 , x 2 ) • ( f (x 2 , x 1 ) -f (x 1 , x 2 ) ) L a s ( f (x 1 , x 2 ) )

.

Moreover, {Z t } t≥0 is reversible and its invariant distribution π has density

dπ(x 1 , x 2 ) ∝ µ(x 1 , x 2 )dx 1 dx 2 (3.4)

with respect to the Lebesgue measure on R 2d , where µ(x 1 , x 2 ) is defined in (2.5).

In Lemma 3.2, the first two terms L a 1 and L a 2 on the right-hand side of (3.3) correspond to the standard Langevin diffusion, while the last term L a s arises from swapping. The replica exchange Langevin diffusion {Z t } t≥0 defined in (2.4)-(2.6) is ergodic, which means that, with any initialization, the distribution of Z t converges to the invariant distribution π. There are two perspectives to characterize the convergence of the Markov process {Z t } t≥0 : (i) the χ 2 -divergence between the distribution of Z t and the invariant distribution π, and (ii) the convergence of the empirical measure of {Z t } t≥0 , which is viewed as a random element in the space of measures. In the following, we quantify the convergence of {Z t } t≥0 from both perspectives. In §3.2, we use Poincaré inequality to quantify (i), while in §3.3, we apply the large deviation principle (LDP) to characterize (ii). In particular, we show that the term L a s in (3.3) plays a crucial role in accelerating the convergence.

3.2. CONVERGENCE IN χ 2 -DIVERGENCE:

Let µ t be the distribution of the replica exchange Langevin diffusion {Z t } t≥0 at time t, and π be its invariant distribution. In the following, we quantify the discrepancy between µ t and π using the χ 2 -divergence, which is defined as

χ 2 (µ t ∥ π) = ∫ ( dµ t dπ -1 ) 2 dπ,

where dµ t /dπ is the Radon-Nikodym derivative between µ t and π. In the following, we characterize the evolution of χ 2 (µ t ∥ π) along time t using the Dirichlet form, which is defined as

E a (f ) = ∫ Γ a (f )dπ, where Γ a (f ) = 1/2 • ( L a (f 2 ) -2f L a (f ) ) . (3.5)

Here Γ a is called Carré du Champ operator and recall that we use π and L a to denote the invariant distribution and infinitesimal generator of {Z t } t≥0 , respectively. In other words, the Dirichlet form is defined as the integration of the Carré du Champ operator under the invariant distribution π.

To characterize the evolution of the χ 2 -divergence, we take the derivative of χ 2 (µ t ∥ π) with respect to time t. Recall that µ t is the distribution of Z t , and by (B.2) in the appendix the corresponding semigroup {P t } t≥0 is defined as

P t (f (x)) = E[f (Z t ) | Z 0 = x].

By setting f as dµ 0 /dπ and the definition of conditional expectation, we have dµ t /dπ = P t (dµ 0 /dπ), which implies

d dt χ 2 (µ t ∥ π) = d dt ∫ [ P t ( dµ 0 dπ )] 2 dπ = 2 ∫ P t ( dµ 0 dπ ) • d dt [ P t ( dµ 0 dπ )] dπ = 2 ∫ P t ( dµ 0 dπ ) • L a ( P t ( dµ 0 dπ )) dπ, (3.6)

where the last equation follows from the definition of the infinitesimal generator in Definition B.2 in the appendix. Meanwhile, by Definition B.3 in the appendix and the invariance of π, we have ∫ L a (f 2 )dπ = 0, which together with (3.5) implies

E a (f ) = - ∫ f L a (f )dπ. (3.7)

Combining (3.6) and (3.7) we obtain

d dt χ 2 (µ t ∥ π) = -2E a ( P t ( dµ 0 dπ )) = -2E a ( dµ t dπ ) , (3.8)

which shows that the derivative of the χ 2 -divergence between µ t and the invariant distribution π is exactly negative twice the Dirichlet form of the Radon-Nykodim derivative dµ t /dπ. In the following theorem, we show that a larger swapping intensity a boosts the Dirichlet form and thus accelerates the convergence. Recall that, as defined in (3.5), E a is the Dirichlet form of the replica exchange Langevin diffusion {Z t } t≥0 with swapping intensity a.

Theorem 3.3. For any fixed function f , E a (f ) is an increasing nonnegative function with respect to a. In particular, we have

EQUATION

where the both terms on the right-hand side are nonnegative.

See §C.2 of the appendix for a detailed proof of Theorem 3.3. Combined with (3.8), Theorem 3.3 shows that swapping accelerates the evolution of χ 2 (µ t ∥ π). In particular, the Dirichlet form E a on the right-hand side of (3.8) decomposes into two terms. The first term E 0 in (3.9) corresponds to the replica exchange Langevin diffusion without swapping, that is, a = 0. More specifically, the two nonnegative terms in E 0 characterize the individual dynamics of the two particles, respectively. In particular, larger temperatures τ 1 and τ 2 yield a larger E 0 , which implies a faster rate of convergence. Meanwhile, as shown in the proof of Theorem 3.3, the specific form of swapping in (2.7) ensures the nonnegativity of the second term in (3.9), which further boosts the Dirichlet form and leads to faster evolution of χ 2 (µ t ∥ π) in (3.8). It is worth mentioning that Theorem 3.3 does not rely on Assumption 3.3 and holds for all objective functions U (•).

To better understand Theorem 3.3, note that the symmetry of f , which corresponds to dµ t /dπ(x 1 , x 2 ) in (3.8), plays a crucial role in the acceleration effect. In specific, when dµ t /dπ(x 1 , x 2 ) is asymmetric, the second term of the Dirichlet form in (3.9) is positive. Otherwise when dµ t /dπ(x 1 , x 2 ) is symmetric, the second term in (3.9) vanishes, since dµ t /dπ(x 1 , x 2 )dµ t /dπ(x 2 , x 1 ) = 0 for all x 1 and x 2 . In other words, the acceleration effect degenerates due to symmetry. Intuitively, at time t the two particles are equivalent and swapping does not change their joint distribution, and hence does not affect the convergence of χ 2 (µ t ∥ π).

Based on (3.8), we further establish the exponential convergence of χ 2 (µ t ∥ π). The key ingredient is the Poincaré inequality (Bakry et al., 2013) . Furthermore, we show that such an exponential rate of convergence is dictated by the constant in the Poincaré inequality, which is called the Poincaré constant. In particular, we show how the swapping intensity a affects the Poincaré constant and accelerates the exponential convergence.

For a Markov process, the Poincaré inequality states that the χ 2 -divergence between any probability measure and its invariant distribution is uniformly upper bounded by the Dirichlet form of the Radon-Nykodim derivative, provided that the probability measure of interest is absolutely continuous with respect to the invariant distribution. In specific, we say that {Z t } t≥0 satisfies the Poincaré inequality with Poincaré constant κ, if for all probability measures ν ≪ π, the following inequality holds,

χ 2 (ν ∥ π) ≤ κ • E a ( dν dπ

) .

(3.10)

Note that the Poincaré inequality is specified by the invariant distribution π and the Dirichlet form E a .

When the Poincaré inequality in (3.10) holds, the derivative of the χ 2 -divergence in (3.8) is upper bounded by itself, which implies the exponential rate of convergence. In particular, we have

d dt χ 2 (µ t ∥ π) ≤ -2κ -1 • χ 2 (µ t ∥ π), which yields χ 2 (µ t ∥ π) ≤ χ 2 (µ 0 ∥ π) • e -2t/κ . (3.11)

Establishing the Poincaré inequality is highly nontrivial. However, for the replica exchange Langevin diffusion, the following theorem shows that Poincaré inequality holds under the smoothness and dissipativity conditions in Assumption 3.1, which implies (3.11). Theorem 3.4. Under Assumption 3.1, the replica exchange Langevin diffusion specified in (2.4)-(2.6) satisfies the Poincaré inequality in (3.10). Hence, the second inequality in (3.11) holds, which implies the exponential decay of χ 2 -divergence.

The proof of Theorem 3.4 adapts from the proof in Bakry et al. (2008) , where a standard Langevin diffusion with unique driving temperature is considered. See §C.3 of the appendix for a detailed proof.

3.3. LARGE DEVIATION PRINCIPLE (LDP) ANALYSIS:

We focus on the empirical measure of the replica exchange Langevin diffusion {Z t } t≥0 defined in (2.4)-(2.6), which is in the same spirit of Polyak averaging (Polyak & Juditsky, 1992) but replaces the iterates with Dirac measures. In specific, for any fixed time T > 0, the empirical measure of {Z t } t≥0 is defined as

ν T = 1 T ∫ T 0 δ Zt dt, (3.12)

where δ Zt denotes the Dirac measure at Z t . Note that {ν T } T ≥0 is a sequence of random measures, which are random elements of P(R 2d ), that is, the space of probability measures on R 2d . We equip P(R 2d ) with the topology of weak convergence, which enables us to define open and closed sets of P(R 2d ).

Recall that Lemma 3.2 shows that the replica exchange Langevin diffusion {Z t } t≥0 is reversible. This fact enables us to apply the Donsker-Varadhan theory (Donsker & Varadhan, 1975) , which implies that {Z t } t≥0 obeys LDP. Formally, we have the following theorem.

Theorem 3.5. The sequence of empirical measures {ν T } T ≥0 of the replica exchange Langevin diffusion {Z t } t≥0 defined in (3.12) obeys LDP. That is to say, for all open sets O and closed sets F in P(R 2d ), the following inequalities hold,

lim inf T →∞ 1 T log P(ν T ∈ O) ≥ -inf ν∈O I a (ν), lim sup T →∞ 1 T log P(ν T ∈ F ) ≤ -inf ν∈F I a (ν).

Here I a (•) : P(R 2d ) → [0, ∞] is the rate function, which takes the form

I a (ν) = ∫ τ 1 • ∇ x1 √ dν/dπ(x 1 , x 2 ) 2 + τ 2 • ∇ x2 √ dν/dπ(x 1 , x 2 ) 2 dπ(x 1 , x 2 ) I 0 (ν) + ∫ a/2 • s(x 1 , x 2 ) • (√ dν/dπ(x 2 , x 1 ) - √ dν/dπ(x 1 , x 2 ) ) 2 dπ(x 1 , x 2 ) Acceleration Effect (3.13)

for all ν ≪ π, and I a (ν) = ∞ otherwise.

See §C.5 of the appendix for a detailed proof of theorem 3.5. The LDP rate function I a (•) in Theorem 3.5 characterizes the rate of convergence for the empirical measures {ν T } T ≥0 towards its population version, which is the invariance distribution of {Z t } t≥0 . More specific, it quantifies the exponential rate of decay of the probability that empirical empirical measures {ν T } T ≥0 deviate from the invariant distribution π. Recall that, under weak topology, P(R 2d ) is metrizable via the Lévy-Prokhorov metric d LP (•, •) (Billingsley, 2013) . We use B r to denote the open ball centered at π with radius r in the metric space (P(R 2d ), d LP ). Then B c r is a closed set and Bc r is open. Then according to Theorem 3.5, we have

lim inf T →∞ 1 T log P ( d LP (ν T , µ) > r ) = lim inf T →∞ 1 T log P(ν T ∈ Bc r ) ≥ -inf ν∈ Bc r I(ν), lim sup T →∞ 1 T log P ( d LP (ν T , µ) ≥ r ) = lim sup T →∞ 1 T log P(ν T ∈ B c r ) ≤ -inf ν∈B c r I(ν).

If we ignore the slight difference between B c r and Bc r , when T is large, the probability P(d LP (ν T , µ) ≥ r) has approximated scale of exp(-T inf ν∈B c r I(ν)). That is to say, inf ν∈B c r I(ν) characterizes the exponential rate of decay of the probability P(d LP (ν T , µ) ≥ r) as T goes to infinity. A larger rate function implies a faster exponential rate of decay. Note that LDP includes both upper and lower bounds in Theorem 3.5 and hence, the exponential rate of convergence characterized by LDP is tight.

To see how swapping accelerates the convergence of the empirical measure towards its invariant distribution, note that the LDP rate function in (3.13) decomposes into two terms. The first term corresponds to the LDP rate function of replica exchange Langevin diffusion without swapping, that is, a = 0. In specific, it is the sum of two nonnegative terms, which characterize the individual dynamics of the two particles, respectively. Meanwhile, the second term in (3.13) is nonnegative and boosts the LDP rate function, which accelerates the convergence of the empirical measure. Similar to the convergence of the χ 2 -divergence in Theorem 3.3, a larger swapping intensity a yields a more significant acceleration effect. The acceleration effect also degenerates when the Radon-Nykodim derivative dν/dπ is symmetric. Finally, it is worth noting that, like Theorem 3.3, Theorem 3.5 does not rely on Assumption 3.3 and holds for all objective functions U (•). In summary, Theorems 3.3 and 3.5 together justify the benefit of swapping from two different perspectives.

Note that the LDP rate function given in Theorem 3.5 takes a different form compared with the one obtained in Dupuis et al. (2012) . Despite the different the forms, the two results are not contradictory. Our formula is derived from the Donsker-Varadhan theory, which uses the assumption of reversibility, while theirs is obtained from a dual formulation of the LDP rate function. Moreover, one can show that they are equivalent.

3.4. DISCRETIZATION ERROR ANALYSIS:

In this section, we lay out the discretization error analysis of replica exchange Langevin diffusion. Our previous discussion in §3.2 and §3.3 focuses on the continuous-time process {Z t } t≥0 defined in (2.4)-(2.6), which needs to be discretized in practice. However, swapping the positions of particles makes it hard to analyze the discretization error, since there is no way to incorporate the swapping of positions into one unified stochastic differential equation. To overcome this difficulty, note that the two particles in the replica exchange Langevin diffusion swap their positions with a specific rate, which is equivalent in distribution to swapping their temperatures with the same rate (Dupuis et al., 2012) . We name the later equivalent process as the temperature swapping Langevin diffusion. With a slight abuse of notations, we still use {Z t } t≥0 to denote the positions of the two particles in the temperature swapping Langevin diffusion. In specific, {Z t } t≥0 is characterized by the following stochastic differential equation,

dZ t = -∇U (Z t )dt + Σ t dW t .

(3.14)

Here Σ t is a random matrix that switches between the diagonal matrices diag{ Note that (3.14) defines a Markov jump diffusion, which can be discretized using the Euler scheme. We denote by (Z (1) (k), Z (2) (k)) and (τ (1) (k), τ (2) (k)) the positions and temperatures of the two particles at discrete time k, respectively. Let (Z (1) (0), Z (2) (0)) = Z 0 and (τ (1) (0), τ (2) (0)) = (τ 1 , τ 2 ) be the initialization. For all integers k ≥ 0, we update the positions of the two particles sequentially as following,

√ 2τ 1 • I d , √ 2τ 2 • I d } and diag{ √ 2τ 2 • I d , √ 2τ 1 • I d },

Z (i) (k + 1) = Z (i) (k) -η • ∇U ( Z (i) (k) ) + √ 2η • τ (i) (k) • ξ (i) (k), i = 1, 2 (3.15)

and swap temperatures according to the following rule,

( τ (1) (k + 1), τ (2) (k + 1) ) = ( τ (2) (k), τ (1) (k) ) with probability a • η • s ( Z (1) (k), Z (2) (k) ) , (3.16) ( τ (1) (k + 1), τ (2) (k + 1) ) = ( τ (1) (k), τ (2) (k) ) with probability 1 -a • η • s ( Z (1) (k), Z (2) (k)

) .

Here {ξ (1) (k)} k≥0 and {ξ (2) (k)} k≥0 are two sequences of independent and identically distributed d-dimensional Gaussian random vectors, η > 0 is the stepsize, while a and s(•, •) are specified in (2.6). Note that η and a are chosen such that η • a < 1 in order to define a valid probability.

Hereafter we denote by {Z η (k)} k≥0 the iterates of the algorithm specified in (3.15) and (3.16). We use the superscript η to emphasize that the iterates depend on the stepsize η. Let {Z η t } t≥0 be the continuous-time interpolation of {Z η (k)} k≥1 , which is a continuous-time stochastic process defined as

Z η t = Z 0 - ∫ t 0 ∇U ( Z η ⌊s/η⌋η ) ds + ∫ t 0 Σ η ⌊s/η⌋η dW s , (3.17)

where

Σ η kη = diag{ √ 2τ (1) (k) • I d , √ 2τ (2) (k) • I d }.

Then for all integers k ≥ 0 and t = kη, we have Z η t = Z η kη = Z η (k). In the following theorem, we characterize the discretization error by the mean squared error between the continuous-time temperature swapping Langevin diffusion in (3.14) and the continuous-time interpolation in (3.17) of the discrete-time algorithm in (3.15) and (3.16). In specific, we consider a fixed time interval [0, T ] and upper bound the mean squared error E[∥Z t -Z η t ∥ 2 ] for all t ∈ [0, T ]. The following theorem shows that this error grows linearly with respect to the stepsize η.

Theorem 3.6. Under Assumption 3.1, there exists a constant γ(d, τ 1 , τ 2 , a, L, α, β, T ) that only depends on the dimension d, temperature parameters τ 1 and τ 2 , swapping intensity a, smoothness constant L and dissipative constants (α, β) of U (•), and length of the time interval T , such that for all t ∈ [0, T ],

E [ ∥Z t -Z η t ∥ 2 ] ≤ γ(d, τ 1 , τ 2 , a, L, α, β, T ) • η,

provided that the stepsize η satisfies 0 < η < α/L 2 .

The proof of Theorem 3.6 adapts from the proof framework of Theorem 5.13 in Yin & Zhu (2010) .

We defer the detailed proof to §C of the appendix.

A DISCUSSION AND CONCLUSION

In this paper, we apply the idea of replica exchange to the Langevin diffusion in the context of nonconvex optimization. In particular, we quantify the benefits of replica exchange compared with the standard Langevin diffusion from two perspectives, that is, the convergence of the χ 2 -divergence and the LDP of the empirical measures. We show from both perspectives that replica exchange accelerates the convergence towards the invariant distribution. In the following, we discuss several related issues.

Flat and sharp minima: The Langevin diffusion yields an invariant distribution that concentrates around the global minima. In particular, as illustrated in Keskar et al. ( 2016), flat minima generally achieves better generalization than sharp ones. Meanwhile, as observed in Zhang et al. (2018) , the invariant distribution of the Langevin diffusion biases more towards the flat minima than the sharp minima. Intuitively, this observation follows from locally integrating the density of the invariant distribution, which is proportional to exp{-U (x)/τ }. As specified in (2.3), the invariant distribution of the replica exchange Langevin diffusion takes the same form, and hence processes the same desired property of biasing towards the flat minima.

Choice of swapping intensity: Recall that in §3 we prove that swapping accelerates the exponential rate of convergence of the χ 2 -divergence and increases the LDP rate function of the empirical measures. In particular, a large swapping intensity a leads to a stronger acceleration effect. In continuous time, the swapping intensity a should be as large as possible. However, in discrete time, to ensure the existence of a valid swapping probability in (3.16), a can only be as large as 1/η, since by (2.6) we have s(•, •) ≤ 1. This constraint characterizes the tradeoff between the acceleration effect and the discretization effort, since more fine-grained discretization with a smaller stepsize η allows for a larger swapping intensity a. Moreover, intuitively speaking, more frequent swapping makes the particles more volatile, which lowers the accuracy of discretization.

Choice of temperatures: In the replica exchange Langevin diffusion, we use two particles driven by low and high temperatures to achieve "local exploitation" and "global exploration", respectively.

There is a fundamental tradeoff in the choice of the temperatures, especially the high temperature.

In specific, a larger τ 2 results in faster global exploration. However, if τ 2 is too large, by the second equation in (2.4), the gradient term ∇U (•) becomes negligible compared with the white noise term √ 2τ 2 dW t . Consequently, the second particle "blindly" explores the whole domain R d uniformly at random, without adapting the geometry of the objective function. As discussed in §2, most swaps happen only when better positions with smaller objective values are discovered by the second particle. Thus, if τ 2 is too large, the global exploration is fast but ineffective.

B BACKGROUND:

In this section, we provide the background of Markov process, which is necessary for the theoretical analysis. We analyze the properties of Markov processes from the viewpoint of Markov semigroup. A semigroup is a family of linear operators on Banach space C(R d ), the space of bounded continuous functions on R d equipped with the uniform norm. Formally, we have the following definition. (1) For t = 0, we have P t = I, which is the identical mapping on C(R d ).

(2) The map t → P t is continuous in the sense that, for all

f ∈ C(R d ), t → P t (f ) is a continuous map from R + to C(R d ).

(3) For all f ∈ C(R d ) and s, t > 0, we have P s+t (f ) = P s (P t (f )).

Moreover, a semigroup of operators is Markov if and only if the following additional conditions are satisfied:

(4) For all t > 0, we have P t (1) = 1, where 1 is the constant function with value one.

(5) For all f ≥ 0, we have P t (f ) ≥ 0.

Recall that a Markov process on R d is a stochastic process {X t } t≥0 satisfying

P(X t ∈ Γ | F s ) = P(X t ∈ Γ | X s ) (B.1)

for all measurable sets Γ and times s < t. Here F s = σ(X u , u ≤ s), which is the natural filtration of the Markov process {X t } t≥0 . Given a Markov process {X t } t≥0 , we define a family of operators {P t } t≥0 on C(R d ), the space of bounded continuous functions on R d , as follows

P t ( f (x) ) = E [ f (X t ) | X 0 = x ] . (B.2)

According to the Markov property in (B.1), we have

P t+s (f ) = P t (P s (f )) (B.3)

and hence, {P t } t≥0 indeed forms a semigroup. We call {P t } t≥0 the Markov semigroup associated with the Markov process {X t } t≥0 .

Given a Markov semigroup, we can define its infinitesimal generator, a linear operator that describes the Markov semigroup's behavior for an infinitesimal t. Formally, we have the following definition.

See, e.g., Revuz & Yor (2013) for more details.

Definition B.2 (Infinitesimal Generator of Markov Semigroup). The infinitesimal generator L of a:

Markov semigroup {P t } t≥0 is defined by

L (f ) = lim t→0 P t (f ) -f t for all f ∈ D(L ).

Here D(L ) denotes the subset of C(R d ) such that the above limit exists.

Intuitively, the infinitesimal generator can be viewed as the derivative of the Markov semigroup at time t = 0. It uniquely determines the Markov semigroup according to the semigroup property (B.3). In the following, we define the invariant and reversible measure with respect to a Markov semigroup.

Definition B.3 (Invariance and Reversibility:

). For a Markov semigroup {P t } t≥0 whose infinitesimal generator is L , a probability measure π is invariant with respect to {P t } t≥0 if and only if

∫ L (f )dπ = 0 for all t ≥ 0 and f ∈ C(R d ). A probability measure π is reversible with respect to {P t } t≥0 if and only if ∫ f L (g)dπ = ∫ gL (f )dπ

for all f, g ∈ D(L ). Note that reversibility implies invariance when we plug g = 1 into the definition of reversibility.

We remark that the standard definitions of invariance and reversibility are based on the Markov semigroup {P t } t≥0 instead of its infinitesimal generator L , which is equivalent to the following definition. The definitions via infinitesimal generator simplify the proof of our lemmas and theorems. In the definitions of Markov semigroup and associated concepts, we restrict the test function f to C(R d ), the space of bounded continuous functions. However, for π being the invariant distribution with respect to {P t } t≥0 , we can always (slightly) extend the definition of Markov semigroup to a larger function space L 2 (π). We use { Pt } t≥0 and L to denote the corresponding extended Markov semigroup and infinitesimal generator. One can show that, when π is reversible, L is a self-adjoint operator in the Hilbert space L 2 (π), which is equipped with the inner product ⟨f, g⟩ π = ∫ f gdπ. Furthermore, one can show that L is positive semidefinite. In this paper, for the convenience of discussion, we do not distinguish the slight difference between L and its extension L , and also the difference of (Phillips & Hille, 1957) . When we choose a test function, we also assume that all the related operations are well-defined.

D(L ), C(R d ), and L 2 (π), since C(R d ) is dense in L 2 (π) and D(L ) is dense in C(R d ) according to Hille-Yosida theorem

C DETAILED PROOFS:

In this section, we present the detailed proofs for all the theorems and lemmas in this paper.

C.1 PROOF OF LEMMA 3.2

Proof. The dynamics in (2.4)-(2.6) defines a Langevin diffusion with jump. Its generator takes the form

L a ( f (x 1 , x 2 ) ) = - ⟨ ∇ x1 f (x 1 , x 2 ), ∇ x1 U (x 1 ) ⟩ + τ 1 ∆ x1 f (x 1 , x 2 ) L a 1 ( f (x 1 , x 2 ) ) - ⟨ ∇ x2 f (x 1 , x 2 ), ∇ x2 U (x 2 ) ⟩ + τ 2 ∆ x2 f (x 1 , x 2 ) L a 2 ( f (x 1 , x 2 ) ) + a • s(x 1 , x 2 ) • ( f (x 2 , x 1 ) -f (x 1 , x 2 ) ) L a s ( f (x 1 , x 2 ) ) (C.1)

and its domain D(L a ) = C 2 c (R 2d ), which is the space of all twice-differentiable functions with compact support. The first two terms L a 1 and L a 2 on the right-hand side of (C.1) correspond to the standard Langevin diffusion, while the last term L a s arises from swapping. To prove invariance and reversibility, by Definition B.3, we only need to show that ∫

R d ∫ R d g(x 1 , x 2 )L a ( f (x 1 , x 2 ) ) dπ(x 1 , x 2 ) = ∫ R d ∫ R d f (x 1 , x 2 )L a ( g(x 1 , x 2 ) ) dπ(x 1 , x 2 ) (C.2)

for all f, g ∈ D(L a ), where π is defined in (3.4).

Note that for the Laplacian term in L a 1 on the right-hand side of (C.1), by integration by parts and the fact that f and g have compact support, we have that for all fixed

x 2 ∈ R d , - ∫ R d g(x 1 , x 2 ) exp ( -U (x 1 )/τ 1 ) ∆ x1 f (x 1 , x 2 )dx 1 (C.3) = ∫ R d ⟨ ∇ x1 [ g(x 1 , x 2 ) exp ( -U (x 1 )/τ 1 )] , ∇ x1 f (x 1 , x 2 ) ⟩ dx 1 = ∫ R d ⟨ exp ( -U (x 1 )/τ 1 ) ∇ x1 g(x 1 , x 2 ) -exp ( -U (x 1 )/τ 1 ) g(x 1 , x 2 )∇ x1 U (x 1 ), ∇ x1 f (x 1 , x 2 ) ⟩ dx 1 .

Recall that µ(x 1 , x 2 ) is defined in (2.5). Hence, (C.3) takes the equivalent form ∫

R d g(x 1 , x 2 )L a 1 ( f (x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 = - ∫ R d ⟨ ∇ x1 f (x 1 , x 2 ), ∇ x1 g(x 1 , x 2 ) ⟩ µ(x 1 , x 2 )dx 1 ,

and hence, ∫

R d ∫ R d g(x 1 , x 2 )L a 1 ( f (x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2 = - ∫ R d ∫ R d ⟨ ∇ x1 f (x 1 , x 2 ), ∇ x1 g(x 1 , x 2 ) ⟩ µ(x 1 , x 2 )dx 1 dx 2 . (C.4)

By switching the positions of f and g in (C.4), we have ∫

R d ∫ R d f (x 1 , x 2 )L a 1 ( g(x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2 = - ∫ R d ∫ R d ⟨ ∇ x1 g(x 1 , x 2 ), ∇ x1 f (x 1 , x 2 ) ⟩ µ(x 1 , x 2 )dx 1 dx 2 . (C.5) By (C.4) and (C.5) we have ∫ R d ∫ R d g(x 1 , x 2 )L a 1 ( f (x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2 = ∫ R d ∫ R d f (x 1 , x 2 )L a 1 ( g(x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2 . (C.6)

By the same derivation, (C.6) also holds for L a 2 . Thus, it remains to prove that (C.6) holds for L a s as well. For notational simplicity, in the following we use f + and f -to denote f (x 1 , x 2 ) and f (x 2 , x 1 ) respectively, and define g + , g -, µ + , and µ -in a similar way. Then we have ∫

R d ∫ R d g + ( 1 ∧ (µ -/µ + ) ) (f --f + )µ + dx 1 dx 2 = ∫ ∫ {(x1,x2):µ -≤µ + } g + µ -(f --f + )dx 1 dx 2 + ∫ ∫ {(x1,x2):µ ->µ + } g + µ + (f --f + )dx 1 dx 2 = ∫ ∫ {(x1,x2):µ -≤µ + } g + µ -(f --f + )dx 1 dx 2 + ∫ ∫ {(x1,x2):µ + >µ -} g -µ -(f + -f -)dx 2 dx 1 = ∫ ∫ {(x1,x2):µ -≤µ + } µ -(g + f -+ g -f + -g + f + -g -f -)dx 1 dx 2 , (C.7)

where the third equality follows from symmetry. Similar to (C.7), we also have ∫

R d ∫ R d f + ( 1 ∧ (µ -/µ + ) )( g --g + ) µ + dx 1 dx 2 = ∫ ∫ {(x1,x2):µ -<µ + } µ -(f + g -+ f -g + -f + g + -f -g -)dx 1 dx 2 . (C.8)

Combining (C.7) and (C.8), we obtain ∫

R d ∫ R d g + ( 1 ∧ (µ -/µ + ) ) (f --f + )µ + dx 1 dx 2 = ∫ R d ∫ R d f + ( 1 ∧ (µ -/µ + ) ) (g --g + )µ + dx 1 dx 2 , which is equivalent to ∫ R d ∫ R d g(x 1 , x 2 )L a s ( f (x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2 = ∫ R d ∫ R d f (x 1 , x 2 )L a s ( g(x 1 , x 2 ) ) µ(x 1 , x 2 )dx 1 dx 2

by the definition of L a s in (C.1). In summary, we obtain that (C.2) holds for π defined in (3.4). Thus, {Z t } t≥0 is a reversible Markov process with the invariant distribution dπ(

x 1 , x 2 ) ∝ exp(-U (x 1 )/τ 1 -U (x 2 )/τ 2 )dx 1 dx 2 .

C.2 PROOF OF THEOREM 3.3

Proof. Recall the infinitesimal generator of the replica exchange Langevin diffusion with swapping intensity a defined in (3.3). According to the properties of ∇ and ∆, for i = 1, 2, we have

∇ xi f 2 (x 1 , x 2 ) = 2f (x 1 , x 2 ) • ∇ xi f (x 1 , x 2 ) ∆ xi f 2 (x 1 , x 2 ) = 2f (x 1 , x 2 ) • ∆ xi f (x 1 , x 2 ) + 2∥∇ xi f (x 1 , x 2 )∥ 2 .

Then by the definition of the Carré du Champ operator Γ a in (3.5), we have

Γ a ( f (x 1 , x 2 ) ) = 1/2 • L a ( f (x 1 , x 2 ) 2 ) -f (x 1 , x 2 ) • L a ( f (x 1 , x 2 ) ) = τ 1 • ∇ x1 f (x 1 , x 2 ) 2 + τ 2 • ∇ x2 f (x 1 , x 2 ) 2 + a/2 • s(x 1 , x 2 ) • ( f (x 2 , x 1 ) -f (x 1 , x 2 ) ) 2 .

Hence, the Dirichlet form E a is given by

E a (f ) = ∫ ( τ 1 • ∇ x1 f (x 1 , x 2 ) 2 + τ 2 • ∇ x2 f (x 1 , x 2 ) 2 ) dπ(x 1 , x 2 ) + ∫ a/2 • s(x 1 , x 2 ) • ( f (x 2 , x 1 ) -f (x 1 , x 2 ) ) 2 dπ(x 1 , x 2 ),

which concludes the proof of Theorem 3.3.

C.3 PROOF OF THEOREM 3.4

Proof. Note that the replica exchange Langevin diffusion without swapping, that is, a = 0, shares the same invariant distribution in (2.5) with the ones with swapping, that is, a > 0. Furthermore, Theorem 3.3 shows that E a (f ) ≥ E 0 (f ) for all a > 0. Hence it remains to show that the replica exchange Langevin diffusion without swapping satisfies the Poincaré inequality in (3.10).

Recall that the replica exchange Langevin diffusion without swapping is defined as

dZ (1) t = -∇U (Z (1) t )dt + √ 2τ 1 dW (1) t , dZ (2) t = -∇U (Z (2) t )dt + √ 2τ 2 dW (2) t , (C.9)

and its infinitesimal generator is given by

L 0 ( f (x 1 , x 2 ) ) = - ⟨ ∇ x1 f (x 1 , x 2 ), ∇ x1 U (x 1 ) ⟩ + τ 1 • ∆ x1 f (x 1 , x 2 ) - ⟨ ∇ x2 f (x 1 , x 2 ), ∇ x2 U (x 2 ) ⟩ + τ 2 • ∆ x2 f (x 1 , x 2 ). (C.10)

The Poincaré inequality for the replica exchange Langevin diffusion in (C.9) is a direct consequence of the existence of a Lyapunov function. In the specific, we have the following lemma, which is adapted from Theorem 1.4 in Bakry et al. (2008) .

Lemma C.1 (Lyapunov implies Poincaré). Let {Z t } t≥0 be the replica exchange Langevin diffusion without swapping on R 2d , whose infinitesimal generator is L 0 defined in (C.10). We further assume that there exists a function V (x 1 , x 2 ) ≥ 1, and constants λ > 0, b ≥ 0, and r > 0 such that

L 0 ( V (x 1 , x 2 ) ) ≤ -λ • V (x 1 , x 2 ) + b • 1 Br (x 1 , x 2 ), (C.11)

where 1 Br denotes the indicator function of the centered ball with radius r in R 2d . Then {Z t } t≥0 satisfies the Poincaré inequality. We call V (x 1 , x 2 ) the Lyapunov function.

Meanwhile, for the first term on the right-hand side of (C.14) we have

∫ ⟨ ∇ x1 [ f 2 u (x 1 , x 2 )/V (x 1 , x 2 ) ] , ∇ x1 V (x 1 , x 2 ) ⟩ • dπ(x 1 , x 2 ) = ∫ ( ∇ x1 f u (x 1 , x 2 ) 2 -∇ x1 f u (x 1 , x 2 ) - ( f u (x 1 , x 2 )/V (x 1 , x 2 ) ) • ∇ x1 V (x 1 , x 2 ) 2 ) • dπ(x 1 , x 2 ) ≤ ∫ ∇ x1 f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ), (C.15)

where the first equality follows from expanding the gradient of f 2 u (x 1 , x 2 )/V (x 1 , x 2 ) with respect to x 1 . By a similar derivation, the same inequality holds for the second term on the right-hand side of (C.14) as well. Hence, plugging them into (C.15), for the first term on the right-hand side of (C.12) we have

∫ -L 0 ( V (x 1 , x 2 ) )/( λV (x 1 , x 2 ) ) • f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) ≤ 1/λ • ∫ ( τ 1 ∇ x1 f u (x 1 , x 2 ) 2 + τ 2 ∇ x2 f u (x 1 , x 2 ) 2 ) • dπ(x 1 , x 2 ). (C.16)

For the second term on the right-hand side of (C.12), we restrict our attention to the bounded domain B r . Based on the assumption that V (x 1 , x 2 ) ≥ 1, there exists a positive constant k such that

∫ b • 1 Br (x 1 , x 2 ) /( λV (x 1 , x 2 ) ) • f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) ≤ b/λ • ∫ Br f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) ≤ k • ∫ Br ∇f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) + (∫ Br f u (x 1 , x 2 ) • dπ(x 1 , x 2 ) ) 2 / π ( B r ) , (C.17)

where the second inequality follows from the Poincaré inequality for measures on a bounded domain (Bakry et al., 2008) . Since (C.17) holds for all functions f u (x 1 , x 2 ), one can choose a suitable u * such that ∫ Br f u * dπ = 0. Then we have

∫ b • 1 Br (x 1 , x 2 ) /( λV (x 1 , x 2 ) ) • f u (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) ≤ k • ∫ Br ∇f u * (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ). (C.18)

Finally, by plugging (C.16) and (C.18) into (C.12), we have that there exists a positive constant ρ such that

Var π (f ) ≤ ∫ f u * (x 1 , x 2 ) 2 • dπ(x 1 , x 2 ) ≤ ρ • ∫ ( τ 1 ∇ x1 f (x 1 , x 2 ) 2 + τ 2 ∇ x2 f (x 1 , x 2 ) 2 ) • dπ(x 1 , x 2 ),

which concludes the proof of Lemma C.1.

C.5 PROOF OF THEOREM 3.5

According to Lemma 3.2, the replica exchange Langevin diffusion {Z t } t≥0 is reversible. Hence, we can apply the Donsker-Varadhan theory, which states that for a continuous-time reversible Markov process with invariant distribution π and infinitesimal generator L , LDP holds and its LDP rate function takes the explicit form

I(ν) = { √ -L (√ dν/dπ ) 2 π , if ν ≪ π, ∞, otherwise. (C.19)

Here √ L denotes the square root of -L , which is defined as follows.

Let A be a positive semidefinite self-adjoint operator in the Hilbert space (H, ⟨•, •⟩ H ), the square root of A is defined as the self-adjoint operator B such that A = B 2 , which means ⟨x, Ax⟩ H = ⟨x, B 2 x⟩ H = ⟨Bx, Bx⟩ H for every x ∈ H. As explained in §B, the infinitesimal generator -L can be extended to a positive semidefinite self-adjoint operator in the Hilbert space L 2 (π). We ignore the slight differences caused by extension. Hence, √ -L is well-defined here.

Recall that the infinitesimal generator L a of the replica exchange Langevin diffusion is given by (3.3). We first derive the explicit form of ∥ √ -L a (f )∥ 2 π . Since the square root of a positive semidefinite self-adjoint operator is self-adjoint, we have

√ -L a (f ) 2 π = ⟨√ -L a (f ), √ -L a (f ) ⟩ π = ⟨ f, √ -L a 2 (f ) ⟩ π = ⟨ f, -L a (f ) ⟩ π = - ∫ f L a (f )dπ.

Since π is the invariant distribution, by Definition B.3 we have ∫ L a (f 2 )dπ = 0. Then we have

√ -L a (f ) 2 π = 1/2 • ∫ ( L a (f 2 ) -2f L a (f ) ) dπ,

which is exactly the Dirichlet form. Then based on (3.3), we obtain

L a (f 2 ) -2f L a (f ) = 2τ 1 • ∇ x1 f (x 1 , x 2 ) 2 + 2τ 2 • ∇ x2 f (x 1 , x 2 ) 2 + a • s(x 1 , x 2 ) • ( f (x 2 , x 1 ) -f (x 1 , x 2 ) ) 2 .

Hence, we have that for probability measures ν ≪ π,

I a (ν) = ∫ τ 1 ∇ x1 √ dν/dπ(x 1 , x 2 ) 2 + τ 2 ∇ x2 √ dν/dπ(x 1 , x 2 ) 2 + a/2 • s(x 1 , x 2 ) • (√ dν/dπ(x 2 , x 1 ) - √ dν/dπ(x 1 , x 2 ) ) 2 dπ, (C.20)

and I(ν) = ∞, otherwise.

C.6 PROOF OF THEOREM 3.6

Proof. Recall that the temperature swapping Langevin diffusion {Z t } t≥0 and the continuous-time interpolated process {Z η t } t≥0 are defined in (3.14) and (3.17). For all t ∈ [0, T ], we have

Z t -Z η t = - ∫ t 0 ( ∇U (Z s ) -∇U ( Z η ⌊s/η⌋η ) ) ds + ∫ t 0 ( Σ s -Σ η ⌊s/η⌋η ) dW s .

By applying the Cauchy-Schwartz inequality and then taking expectation, we have

E [ ∥Z t -Z η t ∥ 2 ] ≤ 2 • E [ ∫ t 0 ( ∇U (Z s ) -∇U ( Z η ⌊s/η⌋η ) ) ds 2 ] + 2 • E [ ∫ t 0 ( Σ s -Σ η ⌊s/η⌋η ) dW s 2 ] . (C.21)

In the sequel, we upper bound the two terms on the right-hand side of (C.21).

For the first term, by Cauchy-Schwarz inequality, we have

E [ ∫ t 0 ( ∇U (Z s ) -∇U ( Z η ⌊s/η⌋η ) ) ds 2 ] ≤ t • E [∫ t 0 ∇U (Z s ) -∇U ( Z η ⌊s/η⌋η ) 2 ds ] ≤ 2L 2 t • ( E [∫ t 0 ∥Z s -Z η s ∥ 2 ds ] + E [∫ t 0 Z η s -Z η ⌊s/η⌋η 2 ds ]) , (C.22)

where the second inequality follows from the L-smoothness of U (•) in Assumption 3.1. In the following, we upper bound the second term on the right-hand side of (C.22). Note that we have

E [∫ t 0 Z η s -Z η ⌊s/η⌋η 2 ds ] ≤ ⌊t/η⌋ ∑ k=0 E [∫ (k+1)η kη Z η s -Z η ⌊s/η⌋η 2 ds ] . (C.23)

For all integers k ≥ 0 and s ∈ [kη, (k + 1)η), based on the definition of {Z η t } t≥0 in (3.17), we have

Z η s -Z η ⌊s/η⌋η 2 = ∥Z η s -Z η kη ∥ 2 = -∇U (Z η kη ) • (s -kη) + Σ η kη ∫ s kη dW u 2 .

By applying the Cauchy-Schwartz inequality again, we obtain

Z η s -Z η ⌊s/η⌋η 2 ≤ 2 • ∇U (Z η kη ) 2 • (s -kη) 2 + 2 • Σ η kη ∫ s kη dW u 2 = 2 • ∇U (Z η kη ) -∇U (x * ) 2 • (s -kη) 2 + 2 • Σ η kη ∫ s kη dW u 2 , (C.24)

where x * is a stationary point of U (•), and by definition ∇U (x * ) = 0. Recall that by Assumption 3.1, U (•) is L-smooth. Based on (C.24), we have

Z η s -Z η ⌊s/η⌋η 2 ≤ 2L 2 • ∥Z η kη -x * ∥ 2 • (s -kη) 2 + 2 • Σ η kη ∫ s kη dW u 2 ≤ 4L 2 • ( ∥Z η kη ∥ 2 + ∥x * ∥ 2 ) • (s -kη) 2 + 2 • Σ η kη ∫ s kη dW u 2 . (C.25)

By integrating the second inequality in (C.25) over the interval [kη, (k +1)η] and taking expectation, then plugging it into the right-hand side of (C.23), we obtain

E [∫ (k+1)η kη Z η s -Z η ⌊s/η⌋η 2 ds ] ≤ 4L 2 • ( sup k≥0 E [ ∥Z η kη ∥ 2 ] + ∥x * ∥ 2 ) • η 3 /3 + 2 • ∫ (k+1)η kη E [ Σ η kη ∫ s kη dW u 2 ] ds. (C.26)

Note that on the right-hand side of (C.26), Σ η kη is a diagonal matrix with diagonal entries √ 2τ 1 or √ 2τ 2 , then by Itô isometry, we have

E [ Σ η kη ∫ s kη dW u 2 ] = 2d ∑ j=1 E [ ( Σ η kη (j) ∫ s kη dW (j) u ) 2 ] ≤ 4dτ 2 (s -kη),

where Σ η kη (j) denotes the j-th diagonal entry of the matrix Σ η kη , and W (j) u is the j-th component of the 2d-dimensional Brownian motion W u . Hence, we have

∫ (k+1)η kη E [ Σ η kη ∫ s kη dW u 2 ] ds ≤ ∫ (k+1)η kη 4dτ 2 (s -kη)ds = 2dτ 2 η 2 , (C.27)

which provides an upper bound for the right-hand side of (C.26). Then by plugging (C.27) into (C.26), we obtain

E [∫ (k+1)η kη Z η s -Z η ⌊s/η⌋η 2 ds ] ≤ 4L 2 • ( sup k≥0 E [ ∥Z η kη ∥ 2 ] + ∥x * ∥ 2 ) • η 3 + 4dτ 2 η 2 . (C.28)

Then by plugging (C.28) into (C.23), we have

E [∫ t 0 Z η s -Z η ⌊s/η⌋η 2 ds ] ≤ (1 + t/η) • ( 4L 2 • ( sup k≥0 E [ ∥Z η kη ∥ 2 ] + ∥x * ∥ 2 ) • η 3 + 4dτ 2 η 2

) .

(C.29)

The following lemma shows that if the discretization stepsize η falls into the interval (0, α/L In other words, the mean squared error of discretization grows linearly with respect to the stepsize η, which concludes the proof of Theorem 3.6.

C.7 PROOF OF LEMMA C.2

Proof. Recall that for all integers k ≥ 0, based on the discretization scheme in (3.15) and (3.16), for i ∈ {1, 2}, we have

Z (i) (k + 1) = Z (i) (k) -η • ∇U ( Z (i) (k) ) + √ 2η • τ (i) (k) • ξ (i) (k),

where ξ (i) (k) is a standard d-dimensional Gaussian random vector and the temperature τ (i) (k) takes value in {τ 1 , τ 2 }. Also note that by the definition of the continuous-time interpolated process {Z η t } t≥0 in (3.17), we have Z(k) = Z η kη . Hence, we have

E [ ∥Z η (k+1)η ∥ 2 ] = E [ Z η kη -η • ∇U (Z η kη ) 2 ] + E [ 2η • τ (1) (k) • ∥ξ (1) (k)∥ 2 ] + E [ 2η • τ (2) (k) • ∥ξ (2) (k)∥ 2 ] + 2 • E [ ⟨ Z η (k+1)η -η • ∇U (Z η kη ), √ 2η • Σ η kη • ξ(k) ⟩ ] . (C.36)

Since ξ(k) is independent of Z η kη and by (3.16) the distribution of τ (k) is only determined by Z η kη , the last term on the right-hand side of (C.36) is zero. Moreover, note that each component of τ (k) only takes value in {τ 1 , τ 2 } and τ 1 < τ 2 . Then by (C.36) we have

E [ ∥Z η (k+1)η ∥ 2 ] ≤ E [ Z η kη -η • ∇U (Z η kη ) 2 ] + 4dητ 2 . (C.37)

Recall that U (•) is (α, β)-dissipative and L-smooth by Assumption 3.1. Then we have

E [ Z η kη -η • ∇U (Z η kη ) 2 ] = E [ ∥Z η kη ∥ 2 ] -2η • E [ ⟨ Z η kη , ∇U (Z η kη ) ⟩ ] + η 2 • E [ ∇U (Z η kη ) 2 ] ≤ (1 -2αη + 2η 2 L 2 ) • E [ ∥Z η kη ∥ 2 ] + 2ηβ + 2η 2 L 2 • ∥x * ∥ 2 ,

where x * is a stationary point of U (•). Hence, based on (C.37), we have

E [ Z η (k+1)η 2 ] ≤ (1 -2αη + 2η 2 L 2 ) • E [ ∥Z η kη ∥ 2 ] + 2(β + dτ 2 )η + 2η 2 L 2 ∥x * ∥ 2 . (C.38)

Note that when η ∈ (0, α/L 2 ), we have 1 -2αη + 2η 

:



Figures and Tables:

Figure1: An illustration of the replica exchange Langevin diffusion. In (a) we illustrate the diffusion process driven by the low temperature, which locally exploits the geometry. The orange dashed lines characterize the evolution of the standard Langevin diffusion, while the blue dashed line denotes the swap, which drives the diffusion process into the neighborhood of another local minimum. In (b) we illustrate the trajectories of a pair of diffusion processes driven by two temperatures. The orange and yellow lines correspond to the low and high temperatures, respectively. The black lines are the contours of the objective function. Note that the lines of the same colors are disjoint due to the swap, which is denoted by the red dashed line. In (c) we illustrate the evolution of the χ 2 -divergence. The upper and lower blue curves correspond to diffusion processes with zero and positive swapping intensities, respectively. The two solid red arrows denote the derivatives of the χ 2 -divergence, and the angle between them characterizes the acceleration effect. In (d) we illustrate the concentration of the empirical measures. The horizontal plane denotes the space of measures, whose center is π, that is, the stationary distribution of the replica exchange Langevin diffusion. The yellow and blue surfaces that center at π characterize the probability densities of the empirical measures, corresponding to zero and positive swapping intensities, respectively. Compared with the blue surface, the yellow one is more concentrated around π, which also characterizes the acceleration effect of swapping.

Definition B.1 (Semigroup of Operators). A family {P t } t≥0 of linear operators on C(R d ) is called a semigroup if and only if it satisfies the following conditions:

where I d denotes the d-dimensional identity matrix. The matrix Σ t characterizes the switching of temperatures, and the rate of switching is same as the one specified in (2.6) and (2.7). Note that the temperature swapping Langevin diffusion in (3.14) is equivalent to the replica exchange Langevin diffusion defined in (2.4)-(2.6) in distribution. Furthermore, the stochastic differential equation in (3.14) enables us to define a continuous-time interpolation of the discrete-time sequence, which is useful in the discretization error analysis. Hence, for the purpose of optimization, we only need to discretize and analyze the temperature swapping Langevin diffusion in (3.14).

Table content: <table/>

Definition B.4 (Equivalent Definitions of Invariance and Reversibility ). For a Markov semigroup {P t } t≥0 , a probability measure π is invariant with respect to {P t } t≥0 if and only if ∫ P Reversibility is another important concept in Markov semigroup theory. A probability measure µ is reversible for {P t } t≥0 if and only if ∫ f P t gdπ = ∫ gP t f dπ for all t ≥ 0 and f, g ∈ C(R d ). Based on the Definition B.2, we can show that this definition is equivalent to Definition B.3.

Table content: <table/>

Lemma C.2. If 0 < η < α/L 2 , there exists a constant δ 1 (d, τ 2 , L, α, β), which depends on the dimension d, the temperature parameter τ 2 , and the smoothness constant L and dissipative constants(α, β) of U (•), such that ] ≤ δ 1 (d, τ 2 , L, α, β).Proof. The proof idea is to show that the sequence {E[∥Z η kη ∥ 2 ]} k≥0 satisfies a contractive inequality based on the discretization scheme defined in (3.15). See §C.7 for a detailed proof.By applying Lemma C.2 to (C.29), we have that there exists a constant δ 2 (d, τ 2 , L, α, β) such that which establishes an upper bound for the first term on the right-hand side of (C.21).It remains to upper bound the second term on the right-hand side of (C.21). According to Itô isometry, we have (j) and Σ η kη (j) are the j-th diagonal entries of Σ s and Σ η kη , respectively. Recall that Σ s and Σ η s are diagonal matrices with all diagonal entires being √ 2τ 1 or √ 2τ 2 . Then for all j, k, and possible realizations of Σ s (j) and Σ η kη (j), we have∫ (k+1)ηRecall that the rate of swapping is specified in (2.6) and (2.7). The conditional probability on the right-hand side of (C.33) satisfies are the first and second components of Z η kη , s(•, •) is defined in (2.6), and o(•) is the little-o notation, which denotes the higher-order term with respect to s -kη. Hence, by combining (C.32)-(C.34), we have that there exists a constant δ 3 (τ 1 , τ 2 , a) such that ∫ (k+1)ηThen based on (C.31), there exists a constant δ 4 (d,τ 1 , τ 2 , a) such that (τ 1 , τ 2 , a) • η 2 ≤ δ 4 (d, τ 1 , τ 2 , a) • t • η. • δ 4 (d, τ 1 , τ 2 , a) • t • η.Hence, by applying the Grönwall's inequality(Dragomir, 2003), we have that there exists a constant γ(d, τ 1 , τ 2 , a, L, α, β, T ) such that for all t ∈ [0, T ],E [ ∥Z t -Z η t ∥ 2 ] ≤ γ(d, τ 1 , τ 2 , a, L, α, β, T ) • η.

Table content: <table><tr><td colspan="14">sup k≥0 Z η s -Z η E [ ∥Z η ⌊s/η⌋η 2 ] ≤ 2d • ( ⌈t/η⌉ + 1 2 ds ] ≤ δ 2 (d, τ 2 , L, α, β) • η. ) • δ 3 (C.35) Finally, by plugging (C.30) and (C.35) into (C.21), we have [∫ t 0 E [ ∫ t 0 ( Σ s -Σ η ⌊s/η⌋η ) dW s 0 kη ∥ 2 E E [ ∥Z t -Z η t ∥ 2 ] ≤ 4L 2 t • (∫ t E [ ∥Z s -Z η ) s ∥ 2 ] ds + δ 2 (d, τ 2 , L, α, β) • η + 2</td></tr><tr><td colspan="14">Then based on (C.22), we obtain the following inequality</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="3">E [ ∫ t 0</td><td colspan="2">( ∇U (Z s ) -∇U ( [∫ t</td><td>(</td><td>Z η ⌊s/η⌋η</td><td>) )</td><td>ds ]</td><td>2</td><td>]</td><td>)</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="3">≤ 2L 2 t •</td><td>E</td><td>∥Z s -Z η s ∥ 2 ds</td><td>+ δ 2 (d, τ 2 , L, α, β) • η</td><td>,</td><td>(C.30)</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>0</td></tr><tr><td colspan="3">E [ ∫ t 0</td><td>(</td><td colspan="9">Σ s -Σ η ⌊s/η⌋η</td><td>)</td><td>dW s</td><td>2</td><td>]</td><td>=</td><td>2d ∑ j=1</td><td>∫ t 0</td><td>E [ (</td><td>Σ s (j) -Σ η ⌊s/η⌋η (j) ) 2</td><td>] ds</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>≤</td><td>2d ∑ j=1</td><td>⌈t/η⌉ ∑ k=0</td><td>∫ (k+1)η kη</td><td>E [ (</td><td>Σ s (j) -Σ η kη (j) ) 2</td><td>]</td><td>ds,</td><td>(C.31)</td></tr><tr><td colspan="2">where Σ s kη</td><td colspan="3">E [ (</td><td colspan="9">Σ s (j) -Σ η kη (j) ) 2</td><td>] ds = 4(</td><td>√</td><td>τ 2 -</td><td>√</td><td>τ 2 ) 2 •</td><td>∫ (k+1)η kη</td><td>P (</td><td>Σ s (j) ̸ = Σ η kη (j) )</td><td>ds. (C.32)</td></tr><tr><td colspan="14">To upper bound the probability on the right-hand side of (C.32), we take the expectation conditioning</td></tr><tr><td colspan="10">Z η k , which yields</td><td/><td/><td/></tr><tr><td colspan="9">∫ (k+1)η kη</td><td>P</td><td>(</td><td colspan="3">Σ s (j) ̸ = Σ η kη (j) )</td><td>ds = E [∫ (k+1)η kη</td><td>P</td><td>(</td><td>Σ s (j) ̸ = Σ η kη (j) Z η kη</td><td>)</td><td>] ds .</td><td>(C.33)</td></tr><tr><td/><td/><td colspan="6">P (</td><td colspan="6">Σ s (j) ̸ = Σ η kη (j) Z η kη</td><td>)</td><td>= a • s (</td><td>Z</td><td>η(1) kη , Z</td><td>η(2) kη</td><td>)</td><td>• (s -kη) + o(s -kη),</td><td>(C.34)</td></tr><tr><td>where Z</td><td colspan="6">η(1) kη , Z</td><td colspan="4">η(2)</td><td/><td/></tr><tr><td>kη</td><td/><td colspan="4">E [ (</td><td colspan="8">Σ s (j) -Σ η kη (j) ) 2</td><td>] ds = 4(</td><td>√</td><td>τ 2 -</td><td>√</td><td>τ 2 ) 2 •</td><td>∫ (k+1)η</td></tr></table>

2 L 2 < 1. Hence, according to (C.38), there exists a constant δ 1 (d, τ 2 , L, α, β) such that

Table content: <table><tr><td>sup k≥0</td><td>E [ ∥Z η kη ∥ 2 ]</td></tr></table>

Peer review (iclr2019):
 This paper gives a theoretical analysis of an interesting statistical physics technique known as replica exchange. The basic idea is that Langevin dynamics at low temperature is slow to converge, and that one could potentially boost the convergence by alternating between low and high temperature. At the extreme one could imagine running in parallel a random search and a gradient descent, and ``teleporting" the gradient descent algorithm whenever the random search algorithm finds a point with better value. This makes a lot of sense and it is nice to see a theoretical analysis of this. The mathematics are sound, but I do not know whether it is an appropriate submission for ICLR. One comment from the math side: it would be interesting (albeit probably difficult) to study kappa in (3.10) as a function of a. In particular at face value it looks like one only benefits from taking a larger, so why not study the limiting behavior of a->infty? What is the limiting value of kappa? Can you perform those calculations in the convex case at least?