The Jarzynski estimator

March 2022

We have a family of potentials VtV_t on Rn\mathbb{R}^n, evolving with time, and their associated Gibbs measures ptp_t, ie

pt(x)=eVt(x)ZtwhereZt=RneVt(x)dx.\begin{aligned} &p_t(x) = \frac{e^{-V_t(x)}}{Z_t}&& \text{where}&& Z_t = \int_{\mathbb{R}^n} e^{-V_t(x)}dx.\end{aligned}

The goal is to compute, at least numerically, the partition functions ZtZ_t or equivalently the free energy Ft=logZtF_t = - \log Z_t. In many problems, the family of potentials (Vt)(V_t) connects an unknown potential V0V_0, for which we would like to compute F0F_0, and another potential V1V_1 for which F1F_1 is known.

An elegant way of estimating these free energies uses the Jarzynski identity, discovered in 1996. It is based on the Langevin dynamics associated with the evolving potential VtV_t:

dXt=Vt(Xt)dt+2dWtdX_t = - \nabla V_t(X_t)dt + \sqrt{2}dW_t

with (Wt)(W_t) a Brownian motion. Upon standard conditions on VtV_t, this system – a standard Stochastic Differential Equation – has a solution. Note that the distribution of XtX_t is not ptp_t in general. However, there is a remarkable identity relating the distribution of XtX_t with ptp_t.

Jarzynski's identity states that

E[e0tV˙s(Xs)F˙sds]=1.\mathbf{E}[e^{-\int_0^t \dot{V}_s(X_s) - \dot{F}_sds}] = 1.

Here and after, the dot in V˙t\dot{V}_t or F˙t\dot{F}_t always represents a derivative with respect to time.

The consequence of this identity is that

logE[e0tV˙s(Xs)ds]=FtF0, \log \mathbf{E}[e^{-\int_0^t \dot{V}_s(X_s)ds}] = F_t - F_0,

giving rise to a natural estimator for the free energy difference FtF0F_t-F_0: one samples many paths from the dynamics (2), then computes the path integral wt=0tV˙s(Xs)dsw_t = \int_0^t \dot{V}_s(X_s)ds along each path, and then average the values of ewte^{-w_t} and take the log.

A proof of Jarzynski's identity


In general, exponential of paths integrals are not easy to study. The classical way is to use the Feynman-Kac representation of PDEs, but this proof seems quite involved; instead, there is a nice trick shown to me by Eric Vanden-Eijnden. The trick is as follows: if ρt(x,y)\rho_t(x,y) is the probability density of the system (Xt,e0tV˙s(Xs)F˙sds)(X_t, e^{-\int_0^t \dot{V}_s(X_s) - \dot{F}_sds}) at tt, then

The proof is then finished, since at=pt=1\int a_t = \int p_t = 1. From now on, we'll set Ut=VtFtU_t = V_t - F_t, so that pt(x)=eUt(x)p_t(x) = e^{-U_t(x)}. Clearly, Ut=Vt\nabla U_t = \nabla V_t.

The augmented system

We set

{dXt=Ut(Xt)dt+2dWtdIt=U˙t(Xt)It \begin{cases}dX_t = - \nabla U_t(X_t)dt + \sqrt{2}dW_t \\ dI_t = -\dot{U}_t(X_t)I_t\end{cases}

subject to the starting condition I0=1I_0=1 and X0p0X_0 \sim p_0. Then clearly,

It=e0tU˙s(Xs)ds. I_t = e^{-\int_0^t \dot{U}_s(X_s)ds}.

We consider the system Zt=(Xt,It)Z_t= (X_t, I_t) as a single SDE dZt=f(t,Zt)dt+σdWtd Z_t = f(t,Z_t)dt + \sigma dW_t with

ft(x,y)=(Ut(x)U˙t(x)y) and σ=diag(2,,2,0).\begin{aligned}&f_t(x,y) = \begin{pmatrix} - \nabla U_t(x) \\ -\dot{U}_t(x)y \end{pmatrix} &&\text{ and }&& \sigma = \mathrm{diag}(\sqrt{2}, \dotsc, \sqrt{2}, 0).\end{aligned}

Let us note ρt(x,y)\rho_t(x,y) the density of (Xt,It)(X_t, I_t), and write the associated Fokker-Planck equation; here and after, \nabla is a gradient and \nabla \cdot is the divergence, ie φ=iiφ\nabla \cdot \varphi = \sum_i \partial_i \varphi. Subscripts denote partial differentiation with respect to some variables.

ρ˙t=[ftρt]+σ22Δρt=(ft)ρtftρt+Δxρt=(xft)ρt(yft)ρtftρt+Δxρt=(Ut)ρt+U˙tρtftρt+Δxρt\begin{aligned}\dot{\rho}_t &= -\nabla \cdot \left[ f_t\rho_t \right] + \frac{\sigma^2}{2}\Delta \rho_t\\ &= - (\nabla \cdot f_t) \rho_t - f_t \cdot \nabla \rho_t + \Delta_x \rho_t \\ &= - (\nabla_x \cdot f_t) \rho_t - (\nabla_y \cdot f_t)\rho_t - f_t \cdot \nabla \rho_t + \Delta_x \rho_t \\ &=- (\nabla \cdot \nabla U_t) \rho_t + \dot{U}_t\rho_t - f_t \cdot \nabla \rho_t + \Delta_x \rho_t\\ \end{aligned}

The marginal integration

Set at(x)=0yρt(x,y)dya_t(x) = \int_0^\infty y \rho_t(x,y)dy, so that E[It]=at(x)dx\mathbf{E}[I_t] = \int a_t(x)dx. Multiplying the last equation by yy, integrating, and swapping integrals and derivatives, we get

a˙t=yρ˙t(x,y)dy=(Ut)at+U˙taty(ftρt)+Δat\begin{aligned}\dot{a}_t &= \int y \dot{\rho}_t(x,y)dy \\ &= - (\nabla \cdot \nabla U_t) a_t + \dot{U}_t a_t - \int y (f_t \cdot \nabla \rho_t) + \Delta a_t \end{aligned}

We have ftρt=UtxρtU˙tyyρtf_t \cdot \nabla \rho_t = - \nabla U_t \cdot \nabla_x \rho_t - \dot{U}_t y \nabla_y \rho_t, hence

a˙t=(Ut)at+U˙tat+Utat+U˙ty2yρt(x,y)dy+Δat.\dot{a}_t = - (\nabla \cdot \nabla U_t) a_t + \dot{U}_t a_t + \nabla U_t \cdot \nabla a_t + \dot{U}_t \int y^2 \partial_y \rho_t(x,y)dy+ \Delta a_t.

An integration by parts shows that y2yρt(x,y)dy=2yρt(x,y)dy=2at\int y^2 \partial_y \rho_t(x,y)dy = -2 \int y \rho_t(x,y)dy = -2a_t, providing a sufficient decay of yy2ρ(x,y)y \to y^2 \rho(x,y) at infinity. We obtain the following equation:

a˙t=(Ut)atU˙tat+Utat+Δat. \dot{a}_t = -(\nabla \cdot \nabla U_t)a_t - \dot{U}_t a_t + \nabla U_t \cdot \nabla a_t + \Delta a_t.

It turns out that the density pt(x)=eUt(x)p_t(x) = e^{-U_t(x)} is the solution of this equation. Indeed, pt=U˙tptp_t = - \dot{U}_t p_t and pt=Utpt\nabla p_t = -\nabla U_t p_t, so that it is easy to check that (11) is satisfied. Moreover, a0(x)=yρ0(x,y)dy=eU0(x)=p0(x)a_0(x)= \int y \rho_0(x,y)dy = e^{-U_0(x)} = p_0(x), so the initial conditions are identical.

To conclude, we obtain

E[It]=at(x)dx=pt(x)dx=1.\mathbf{E}[I_t] = \int a_t(x)dx = \int p_t(x)dx = 1.