Accelerated Jarzynski estimator with deterministic virtual trajectories

The Jarzynski estimator is a powerful tool that uses nonequilibrium statistical physics to numerically obtain partition functions of probability distributions. The estimator reconstructs partition functions with trajectories of the simulated Langevin dynamics through the Jarzynski equality. However, the original estimator suffers from slow convergence because it depends on rare trajectories of stochastic dynamics. In this paper, we present a method to significantly accelerate the convergence by introducing deterministic virtual trajectories generated in augmented state space under the Hamiltonian dynamics. We theoretically show that our approach achieves second-order acceleration compared to a naive estimator with the Langevin dynamics and zero variance estimation on harmonic potentials. We also present numerical experiments on three multimodal distributions and a practical example in which the proposed method outperforms the conventional method, and we provide theoretical explanations.

Figure 1

Illustrations of (a) the proposed HJE method and (b) the conventional LJE method for estimating the partition function of a target distribution. Both methods generate trajectories, which are depicted by solid curves, whose initial positions are given by some initial distribution. A single trajectory for each method is colored in red (dark gray) for visibility. A protocol for each process is defined during time $t\in [0,\tau ]$, and boundary values of the protocol are determined by the initial distribution and the target distribution. The main feature of the HJE is that its trajectories are generated by the Hamiltonian dynamics with virtual momentum, which is deterministic given the initial states, while the LJE adopts the Langevin dynamics, which is stochastic.

Figure 2

Analytically obtained total dissipative work $W_{\mathrm{diss}}$ of the HJE and LJE as a function of duration $\tau$ in linear systems. $W_{\mathrm{diss}}$ of the HJE and LJE are depicted by a red solid curve and a blue dashed curve, respectively. (a) Parallel transport of a harmonic potential, where ${\sigma }^2=1$, ${\mu }_{\tau }=1$, and $m_0=1$. (b) Scaling of a harmonic potential, where ${\sigma }_0^2=1$, ${\sigma }_{\tau }^2=4$, and $m\left(t\right)=1$. (c) Scaling of a harmonic potential with time-dependent mass for the HJE, where ${\sigma }_0^2=1$, ${\sigma }_{\tau }^2=4$, and $\alpha =1$. (d) Rotation of a harmonic potential, where $a=8$, $b=2$, and $m_0=1$. As a whole, the specific parameters do not change the results qualitatively.

Figure 3

Total dissipative work $W_{\mathrm{diss}}$ exerted from a harmonic potential that is dragged and scaled taking duration $\tau =2\pi$, where the HJE and LJE are depicted by red (dark gray) and blue (light gray) surfaces, respectively. $\mu$ is the dragged distance and $a$ characterizes how much the potential is scaled.

Figure 4

Error of the estimated partition function for the HJE and LJE, which are depicted by a red solid curve with dots and a blue dashed curve with dots, respectively, as a function of $\tau$ for a 1D double-well potential. The target distribution is characterized by $k=1/16$ and its partition function is $Z=146.372$, which is numerically confirmed. We generated 10 trajectories at each $\tau$ to estimate the partition function and repeated the procedure 100 000 times for both the HJE and LJE. The adopted numeric schemes are the Runge-Kutta fourth-order method for the HJE and the Euler-Maruyama method [30] for the LJE. Applying a higher-order scheme for stochastic differential equations, the order 1.0 strong stochastic Runge-Kutta method [31], for the LJE makes no qualitative difference to the result (not shown). We set the time step for numerical integration to be $\mathrm{\Delta }t={10}^{-3}$.

Figure 5

Illustration of the target distribution and error of the estimated partition function. (a) Plot of the target distribution. The distribution is a mixture of Gaussian distributions, and the peaks are located at points $P=(a,s)$ and $Q=(a,-s)$. We define the displacement $s$ given by the distance between the $x$-axis and the peaks of the target distribution so as to characterize their arrangement. (b) Error for the HJE and LJE, which are depicted by a red solid curve with dots and a blue dashed curve with dots, respectively, as a function of $s$, where $\tau =2\pi$, ${\sigma }^2=1$, and $a=10$. We generated 10 trajectories at each $s$ to estimate the partition function and repeated the procedure 1000 times for the HJE and LJE. The settings for the numerical schemes are the same as in Sec. 4a.

Figure 6

Error of the estimated partition function for the HJE and LJE, which are depicted by red dots and blue circles, respectively, as a function of $N$ characterizing multimodal distributions in 16-dimensional space. Points at each $N$ correspond to different sets of $\left\{{\mu }_i\right\}$. The ground truth of $Z$ is numerically obtained as in Appendix pp9. We tested 20 sets of $\left\{{\mu }_i\right\}$. Then, we generated 1000 trajectories for each set to estimate the partition function, and we repeated the procedure 10 times for the HJE and LJE. The duration of the process is $\tau =4\pi$. The settings for the numerical schemes are the same as in Sec. 4a.

Figure 7

Illustration of the stretched Rouse model and the error of the estimated partition function. (a) Model description. A Rouse model consists of $N$ beads of positions $q_1$ to $q_N$ and harmonic springs connecting them with spring constants $k_2$ to $k_N$. Each end point of the chain is held by optical laser traps with positions $q_0$ and $q_{N+1}$ and spring constants $k_0$ and $k_{N+1}$, respectively. The right trap moves at a constant velocity $v=\mu /\tau$. (b) Error of the estimated partition function with the protocol of the stretching Rouse model for the HJE and LJE, depicted by a red solid curve with dots and a blue dashed curve with dots, respectively, as a function of $\tau$. The model is characterized by $N=100$, $\mu =20$, and $k_n=1$ for $n=1,\cdots ,N$. We generated 10 trajectories at each $\tau$ to estimate the partition function and repeated the procedure 2000 times for both the HJE and LJE. The settings for the numerical schemes are the same as in Sec. 4a.