Integrated tempering enhanced sampling method as the infinite switching limit of simulated tempering

Molecular dynamics simulation is a powerful method for investigating microscopic biochemical systems. However, when the system contains barriers due to energy or entropy, direct molecular dynamics simulation will lead to a very high computational cost. Therefore, people turn to use the enhanced sampling method, sacrificing the real dynamic information in order to capture the desired Gibbs distribution. Until now, sampling methods have already been successfully applied in multiple disciplines including statistical physics, chemical physics, Bayesian statistics, machine learning, and related areas.

The efficiency of the sampling approach highly depends on the convergence rate to the thermodynamic equilibrium. Tempering schemes, such as simulated tempering molecular dynamics (STMD)^1,2 and replica exchange molecular dynamics (REMD),^3–5 were developed and are among the most popular methods to overcome the metastability and to enhance the convergence to equilibrium, thanks to their effectiveness and simplicity to use. The basic idea of both is to use one or more artificial high temperatures to accelerate exploration of the conformational space, while using the physical temperature to sample the desired physical observable. The interaction between different temperatures is designed so as to guarantee the unbiasedness of the estimation. These two methods are quite comparable to each other. STMD gives a higher rate of delivering the system between different temperature states as well a higher rate of traversing the energy space,⁶ but requires the estimation of the partition function first.

Another tempering algorithm named “integrated tempering enhanced sampling” (ITS), which was recently introduced by Gao,⁷ uses a temperature-biased effective potential energy to run the molecular dynamics (MD) simulation and has been successfully applied to various biochemical examples.^8,9 Noting that ITS also uses auxiliary temperatures to accelerate the convergence, it is desirable to compare it with previous methods, and find out why it performs better in many cases, which has not been studied theoretically to the best of our knowledge.

As one of the main results in this paper, we discover that the integrated tempering enhanced sampling method is in fact a reformulated version of the infinite switching limit of STMD, that is, the limit of the simulated tempering when the attempt switching frequency goes to infinity. Moreover, we show that as the switching rate in STMD increases, the empirical measure converges faster toward stationary distribution, using the large deviation principle; and thus the sampling efficiency of STMD increases as the switching frequency increases. Combining the two theoretical findings, we justify the efficiency of ITS over conventional STMD. Finally, we compare the infinite switching STMD against the normal STMD with two numerical examples: an artificial high dimensional system and a more realistic Lennard-Jones example¹⁰ with 16 atoms. The numerical results validate our theoretical findings.

Our study of the switching rate of the STMD is closely related to the recent progress in understanding the swapping frequency of tempering schemes, started with replica exchange methods. For REMD, it has been discovered through numerical examples that the sampling efficiency increases when the frequency of swapping of temperatures is pushed up to infinity;¹¹ but directly increasing the swapping rate to reach this limit is computationally infeasible, as many swaps are needed per MD step. A breakthrough was made in Refs. 12 and 13 which proposed an explicit way to reach this limit and also proved that indeed the sampling efficiency increases in such a limit. A natural reformulation of infinite swapping REMD was later proposed in Ref. 14 which leads to an easy implementation as a simple patch to conventional molecular dynamics. For multiple temperatures, an efficient implementation based on ideas of a multiscale integrator was also proposed for infinite swapping REMD,^15,16 which leads to practical applications in sampling configurational space of large biomolecules. We note that increasing the switching rate in the simulated tempering has been mentioned in Refs. 16 and 17 under a general framework of viewing tempering schemes as MD process augmented by jumping processes, without much details, in particular, the connection with ITS.

In the remaining article, we will first recall the simulated tempering molecular dynamics as a stochastic process in Sec. II. We will justify using a large switching frequency via large deviation principle for the empirical distribution in Sec. III with detailed derivation given in the  Appendix. Infinite switching limit of STMD and its reformulation are derived and generalized in Secs. IV and V. The identification of ITS and infinite switching STMD is given and discussed in Sec. VI, followed by numerical examples in Sec. VII. Some conclusive remarks are given in Sec. VIII.

We begin by recalling the simulated tempering molecular dynamics. For the sake of simplicity, we will first discuss the algorithm for a system governed by overdamped Langevin equations (i.e., when inertia can be neglected). More general dynamics will be discussed in Sec. V. Consider

where $x∈R3n$ denotes the configuration of the system with n particles, f(x) = −∇V(x) is the force associated with the potential V(x), β = 1/k_BT is the inverse temperature, and η is a 3n-dimensional white-noise with independent components. We choose a unit system so that the friction coefficient becomes one for notational simplicity. Equation (1) is consistent with the Boltzmann equilibrium probability density

where the normalization constant $Zβ=∫R3ne−βV(x)dx$ is the partition function.

Assuming we only take two temperatures in simulated tempering (the generalization to more temperatures will be considered in Sec. V), we replace (1) by

where the inverse temperature β(t) now attempts switches between the physical and the artificial temperatures with frequency ν, and the attempted switch from β₁ to β₂ are accepted with the probability

where n₁ and n₂ are some weighting parameters that will be further discussed below.

Note that, in the simulated tempering overdamped dynamics, both the configuration x and the inverse temperature β are dynamical variables: x follows (3), β follows a jump process, and the two dynamics are coupled. To understand the dynamics, it is useful to consider its infinitesimal generator, given by (we use β′ to denote the temperature other than β, namely, β′ = β₂ if β = β₁ and vice versa)

for a smooth function $u:R3n×{β1,β2}→R$⁠. This means that the evolution of a physical observable u under the dynamics is given by

or equivalently, the density ρ of (x, β) evolves under the adjoint operator

For the infinitesimal generator, the first two terms on the right-hand side of (5) correspond to the overdamped dynamics (3), while the last two terms correspond to the jump process of β, − νg_β,β′(x) being the rate of switching from temperature β to β′, and νg_β′,β(x) is the rate of switching to β from the other temperature β′ (attempt switching frequency adjusted by the acceptance rate).

It is straightforward to check the equilibrium distribution of the coupled dynamics of (x, β), as the stationary solution of (7) is given by

where $δβ,β1$ denotes the Kronecker delta: $δβ,β1=1$ if and only if β = β₁. This is nothing but a weighted average of Boltzmann densities at temperatures β₁ and β₂, with weight $n1Zβ1$ and $n2Zβ2$⁠, respectively. On each fixed temperature, the distribution is generated by the ordinary molecular dynamics and thus follows Boltzmann distribution; meanwhile, at each fixed configuration, the proportion of the two temperatures is given by

as determined by the acceptance probability (4). Therefore, (3) together with the jumping process of β samples the ϱ(x, β) as the equilibrium distribution.

Summing ϱ(x, β) over β then gives the marginal equilibrium density for the configuration position alone,

which is a weighted average of the Boltzmann densities at the two temperatures β₁ and β₂. As a result the ensemble average at the physical temperature (β₁) of any observable A(x) can be estimated from

where we have used that the ensemble average equals to the time average thanks to ergodicity. Note that in principle, we shall use two independent realizations to estimate the numerator and denominator on the right-hand side of (11) to get an unbiased estimator, while in practice the bias is well-controlled and the variance of the estimator usually dominates the error.

Note that the above holds for arbitrary positive weighting factors n₁ and n₂, while the sampling efficiency of the method depends on the choice. The conventional wisdom in simulated tempering method is to choose n₁ and n₂ according to the partition function,

With this choice, we have

and thus the sampler spends equal amount time at the two temperatures. It is also possible to choose different n_i to emphasize one of the temperatures. Of course, the partition functions $Zβi$ are not known a priori, and hence the usual approach is to start with some initial guess of the weighting factors and then adaptively adjust them on-the-fly using an iterative method; see, e.g., Refs. 7 and 17–19. In what follows, we will mainly focus on the choice of the switching frequency ν, and thus for the simplicity of discussion, we will assume that n_i are fixed during the sampling process and only make some remarks about adjusting n_i afterwards.

As we have discussed above, for any choice of the switching frequency ν, the trajectory (x^ν, β^ν) (we put superscript ν to emphasize the ν dependence) of the simulated tempering overdamped dynamics can be used to sample the equilibrium distribution (8). In particular, the empirical distribution of the trajectory, defined blow,

will converge to the equilibrium distribution ϱ as T → ∞. Note here δ_(x,β) denotes the Dirac delta function, that is, the point mass function at (x, β), and B is any fixed burn-in time (which for simplicity will be taken as B = 0 in the numerical experiments).

To discuss the choice of ν, it is important then to quantify the speed of convergence of the empirical distribution (14) to the equilibrium: A better ν will correspond to faster convergence and thus less simulation length of the trajectory. To quantify this convergence for STMD, the theoretical tool we will use is the large deviation principle for the empirical measure of stochastic processes. We will discuss the idea and conclusion of the large deviation principle below, while defer the rigorous definition and derivation to the  Appendix.

When T becomes large, the empirical distribution is expected to be very close to the equilibrium. The large deviation principle quantifies the probability that the empirical distribution is still far away from the equilibrium: more specifically, let the probability of the empirical distribution being equal to μ is on the order of $O(e−Iν(μ)T)$⁠, with a rate functional, specific form given below, I^ν(μ) ≥ 0, and only vanish when μ is the equilibrium distribution ϱ. This, in particular, tells us that as T → ∞, the likelihood that empirical distribution is deviated away from the equilibrium is exponentially small, for which the functional I^ν quantifies the rate of the exponential decay. Hence a larger rate function indicates faster rate of convergence.

To specify the rate functional, let μ be a probability measure on $R3n×{β1,β2}$ with a smooth density and define θ(x, β) ≔ [dμ/dϱ](x, β), the ratio of the probability density of μ and the equilibrium distribution. For the simulated tempering process with switching frequency ν, the large deviation rate function for the empirical measure converging to the stationary one is given by

where

Note that J₁ is non-negative and hence the large deviation rate functional is a pointwise monotonic function in ν. Thus, we conclude from this that a larger swapping rate ν corresponds to a faster convergence of the empirical distribution to equilibrium.

As shown in Sec. III, the higher the switching rate ν is, the faster the convergence of the empirical distribution and thus the sampling of the simulated tempering overdamped dynamics is. On the other hand however, a large ν value requires one to make many swapping attempts, which slows down the actual simulation. To resolve this problem, the key is that the limit when ν → ∞ can be taken explicitly. This observation is first made for replica exchange dynamics in Ref. 13. As for simulated tempering, when the switching frequency between the temperatures ν → ∞, the dynamics (3) converges to the following:

where the effective temperature $T(x)=β1−1ω1(x)+β2−1ω2(x)$ is a weighted average of $T1=β1−1$ and $T2=β2−1$⁠. The weight as a function of x is given by (9). Intuitively, this can be understood as with fixed configuration x, the proportion of time the trajectory spends at temperature β_j is given by ω_j(x), and thus the effective temperature is given by a weighted average of the two temperatures.

The Fokker-Planck equation for the probability density of x(t) corresponding to (18) can be written as

where we have defined the effective potential

and the mobility

It is easy to check that the stationary solution of (19) is exp(−β₁U(x)), which is just ϱ(x) by definition of U.

Note that this point of view also leads to a further simplification of (18), in the spirit of Lu and Vanden-Eijnden.¹⁴ We may replace $B$ in the Fokker-Planck equation (19) by a constant function, which, for convenience, we will simply take to be the identity. This substitution does not affect the stationary distribution, and hence preserves the sampling property, but it changes the overdamped Langevin dynamics associated with the Fokker-Planck equation. It is easy to see that this new dynamics is given by

Compare with (18), the noise term is additive in (22), and it still samples the desired stationary distribution ϱ. Note that (22) is very similar to the original overdamped Langevin dynamics (1). The only change is a scaling factor in front of the forcing term, which involves the auxiliary (inverse) temperature β₂ and also the weights ω_j, j = 1, 2. Note that in practice, (22) can be implemented as an easy patch of existing codes for molecular dynamics.

We refer to this infinite switching simulated tempering method as the ITS method.

First, the idea in Sec. IV can be extended naturally to more than two temperatures. For which, the dynamics in the infinite switching limit is given by

where we define the weight here as

if N is the number of temperatures used. Similarly, the dynamics (23) can be viewed as the overdamped dynamics with the effective potential given by (20) where the corresponding marginal equilibrium density becomes

as we have

Thus the gradient of the effective potential U(x) is exactly the forcing term in (23).

The infinite switching simulated tempering can be also generalized to the (underdamped) Langevin equation, rather than the overdamped Langevin dynamics (1); other thermostats can be also used. Recall the Langevin equation

where m denotes the mass and γ the friction coefficient. The generalization of (22) reads

The structure of the equation is rather similar to the overdamped case (22); the only modification compared to (27) is the scaling factor in front of the forcing term that amounts to a weighted average of the inverse temperatures.

The infinite switching limit of the simulated tempering sampling scheme is very closely related to the integrated tempering enhanced sampling (ITS) algorithm originally proposed in Ref. 7. The ITS algorithm introduces a temperature biased effective potential energy as

and runs molecular dynamics simulation on the surface. Here $nβk$ is chosen to be some weighting factor for the temperature β_k. Note that this is identically the same as we have in (20) with the marginal equilibrium ϱ(x) given by (25), despite a constant difference which does not matter in sense of potential energy.

As a conclusion, the ITS algorithm can be viewed as the infinite switching limit of the simulated tempering algorithm. As we discussed in Sec. III, the sampling efficiency of the simulated tempering method increases as ν → ∞. Thus as a corollary, the ITS algorithm is more efficient in sampling compared with the simulated tempering algorithm at a finite switching rate. This will be further demonstrated in our numerical examples in Sec. VII.

We first consider a system in D dimension moving on the following potential with x = (x₀, …, x_D−1):

where λ₁, λ₂, …, λ_D−1 are parameters controlling the stiffness of the harmonic potential in the x₀, x₁, …, x_D−1 directions. We take the low (physical) temperature β₀ = 25 and five artificial temperatures β_k = 25 × 2^−k, k = 1, 2, 3, 4, 5. For each β_k, its weighting parameter n_k are set to be inverse of the corresponding partition function, i.e., $Zβk−1$⁠, which could vary when the number of dimensions changes. When simulating with STMD, we use

so that its infinite swapping limit coincides with (23) [rather than a multi-temperature version of (18) for the original overdamped dynamics (3)]. At a finite switching rate, we just consider the switching attempts of β(t) from some β_k to its adjacent β_k−1 (if k > 0), or β_k+1 (if k < 5). And in each attempt to switch, we shall first decide whether to switch up or down with equal probability, if 0 < k < 5. The total simulation time is T_tot = 2.5 × 10⁶ with time step dt = 0.025, which means the total number of steps is N_tot = 10⁸. As mentioned previously, we will take no burn-in period when estimating the average of physical observable.

To compare the algorithms, we calculate the asymptotic variance of the observable V(x) using a batch estimation,

where AV stands for the asymptotic variance and WS stands for the window size of the batch. The results are shown in Figs. 1 and 2 for D = 1 and D = 10 respectively. We observe that for the finite switching rate, STMD at frequency ν = 1 has a lower asymptotic variance compared to ν = 0.1, and moreover the infinite swapping limit has an even lower asymptotic variance. We also observe that for this example, the simulated tempering with ν = 1 is already quite close to the infinite switching limit.

Now, to test the performance of our algorithm on a more realistic example, we apply (28) for a dimer in a solvent model as considered in Ref. 10. This system consists of N two-dimensional particles in a periodic box with side length l. All particles have the same mass m, and they interact with each other with the Weeks-Chandler-Anderson potential defined as

if r ≤ r_WCA = 2^1/6σ and V_WCA(r) = 0 otherwise, except for a pair of particles which interact via a double well potential

We take N = 16, l = 4.4, σ = 1, h = 1, ω = 0.5, and ϵ = 1 in the simulation. The physical temperature is T₀ = 0.2, and one artificial temperature T₁ = 1.0. The total simulation time is set to be T_tot = 1.0 × 10⁵ with time step dt = 0.001, which means that the total number of steps is N_tot = 10⁸. The quantity of interest is the free energy associated with the distance of the pair of particles interacting via the double well potential.

In this numerical example, we would still choose the weighting parameter to be inverse partition function, which is however not explicitly known or easily obtained a priori in this case. Therefore, we use the following method to get approximation of these partition functions, similar to that of Ref. 7, and then use the estimates in the weighting factors n_k. First, a set of initial guess $Zk(0)$ is chosen and used to run the dynamics simulation. Then from a relatively short trajectory, we can add up the corresponding weight to get the “proportion” of each temperature by normalizing their summation to one. Since this proportion would go to 1/2 if the guess is accurate and the trajectory is infinitely long, we can adjust our guess accordingly. More specifically, if the proportions are $wk$⁠, k = 0, 1, the weighting factor will be updated as

where I is a small neighborhood around 1 to take into account of the fluctuation due to the short simulation trajectory when estimating $wk$⁠. We iteratively update the partition function estimate ${Zk(l)},l=1,2,⋯$ until the proportions $wk$ are satisfactorily close to 1/2. Here we take I = [0.35, 1.5], initial guess $(Z0(0),Z1(0))=(1,108)$⁠, and the number of iterations l_max = 10 with the trajectory length 10⁷ steps each.

The performance of STMD with various switching rates is shown Fig. 3. It can be seen that the performance of STMD is better when the switching rate increases from 0.25 to 25, and the asymptotic variance converges to that of the IST as ν → ∞. This provides strong numerical validation of our theoretical results. We also remark that as this example is more complicated than the previous one, the switching rate needed by STMD to get close to IST becomes much higher, compared to the results in Fig. 2. It is expected that for a larger system and when more temperatures are used, it will become more difficult for STMD to reach its infinite switching limit by increasing the switching rate. Thus it is more advantageous to take the limit explicitly and operate at the infinite switching rate by using, e.g., the IST algorithm.

We justify that the sampling efficiency of the simulated tempering method increases with the switching rate of the temperature, using both the theoretical analysis based on large deviation of empirical distribution and also numerical tests on two examples. This motivates taking the infinite switching limit of the simulated tempering dynamics, which recovers the integrated tempering enhanced sampling method under a natural reformulation. Now the limiting dynamics can be implemented as a patch to standard molecular dynamics by adding a scaling factor to the force term based on a weighted average of all temperatures involved. This leads to a practical scheme with higher sampling efficiency than standard simulated tempering. Throughout this work, we have assumed that the weight factors n_k are given and only consider the efficiency of the algorithm as the parameter ν varies. For practical use of simulate tempering algorithms, it is also essential for us to choose an appropriate set of weights for each temperature or to design an adaptive algorithm to learn a set of optimal weights. A thorough understanding of these issues is beyond the scope of the current paper and is an important direction for future investigation.

H.G. is supported in part by the National Science Foundation of China via Grant No. 11622101. J.L. is supported in part by the National Science Foundation via Grant No. DMS-1454939. J.L. and H.G. would also like to thank Yiqin Gao, Zhiqiang Tan, Eric Vanden-Eijnden, and Zhennan Zhou for helpful discussions.

To introduce formally the large deviation principle, let us introduce some definitions. Let S be a Polish space and $P(S)$ the space of probability measures on S, equipped with the topology of weak convergence. Under this weak topology, $P(S)$ itself is a Polish space. Note that the empirical measure $λTν$ is a $P(S)$-valued random variable.

The large deviation principle for empirical distribution can be stated as follows:^20–22 A sequence of random probability measures {γ_T} is said to satisfy a large deviation principle with a rate function $I:P(S)→[0,∞]$⁠, if for all open sets $O⊂P(S)$

for all closed sets $C⊂P(S)$

and for all M < ∞, {μ: I(μ) ≤ M} is compact in $P(S)$⁠. In particular, once we obtain the rate functional, (A1) and (A2) quantifies how unlikely the empirical distribution is far from the equilibrium when T is large, which was used in this work to justify the infinite switching limit.

To find the rate functional I^ν for the simulate tempering overdamped dynamics with switching rate ν, we divide the infinitesimal generator into two parts,

where

The rate functional thus has an additive structure (15) with the rate functional J₀ and J₁ corresponding to $Ldiffν$ and $Ljumpν$⁠, respectively.

To get an explicit formula of the rate functionals, we first consider J₀. Let $f(x,β)=θ(x,β)$⁠, then following Donsker and Varadhan,²⁰ we have

Next we consider the rate functional J₁ corresponding to the jump process of the temperature with generator $Ljumpν$⁠. From the definition of θ(x, β), we have

and therefore

Combined with the Cauchy-Schwartz inequality, we obtain that for any choice of non-negative u(x, β),

where the equality is attained if and only if for each $x∈R3n$ and each β it holds

In particular, this gives the choice of u(x, β) to make equality holds in (A8).

Using the variational characterization of the large deviation rate functional as in Ref. 20, we have

Therefore, we arrive at the explicit formula for the large deviation rate functionals. We remark that a rigorous proof of the rate functional derived from the above calculation can be found in Ref. 23.