Importance sampling ⚖️⚖️ : The Jarzynski connection

Let $V_t$ be a family of smooth potentials on $\mathbb{R}^n$, continuously evolving with time $t\in[0,1]$. Let $p_t$ be their associated Gibbs measure: 

It is not easy to find a continuous dynamic $X_t$, such that at each time $t$ the distribution of $X_t$ matches $p_t$. One might be tempted to use the overderdamped Langevin dynamics,

with $(W_t)$ a Brownian motion. But the density $\rho_t$ of $X_t$ is not equal to $p_t$ in general. Indeed, $\rho_t$ solves the Fokker-Planck equation

If $\rho_t$ was equal to $p_t$, then $p_t$ should also satisfy this equation; but for $p_t$ one easily checks that the RHS is zero, resulting in $\dot{p}_t=0$ which is false in general if we choose a family of evolving potentials; it would only be true at equilibrium when $V_t$ is constant.

The Jarczynski trick allows to reconnect the distribution of $X_t$ with $p_t$. It introduces a weight, the solution of an auxiliary ODE, such that its conditional expectation given $X_t=x$ gives $p_t(x)$, hence correcting the mismatch between $\rho_t$ and $p_t$.

(In the sequel we note $F_t = \log Z_t$, and the dot in $\dot{V}_t$ or $\dot{F}_t$ always represents a derivative with respect to time. )

Jarczynski's goal was to compute, at least numerically, the partition functions $Z_t$ or equivalently the free energy $F_t = - \log Z_t$. The Jarczynski estimator, discovered in 1996, simply consists in applying (5) to $\varphi=1$, and noting that $W_t$ is a simple ODE solved by

The consequence of this identity is that

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

In general, exponential of paths integrals are not easy to study. The classical way to prove Jarczynski's identity directly is to use the Feynman-Kac representation of PDEs, but this proof seems quite involved; instead, the trick above with the augmented system seems to be easier. It was shown to me by Eric Vanden-Eijnden.

Proof

Equation (5) can be reformulated as follows: if we note $\varrho_t(x,w)$ the density of $(X_t, W_t)$, then

Note that the integral is only on the positive axis because $W_t$ is always positive from (6). We note $q_t(x)$ the term on the right; we are going to prove $p_t=q_t$ by proving that both satisfy the same parabolic equation, then invoke a uniqueness result.

From now on, we'll set $U_t = V_t - F_t$, so that $p_t(x) = e^{-U_t(x)}$ is normalized. Clearly, $\nabla U_t = \nabla V_t$.

The augmented system

We consider the system $Z_t= (X_t, W_t)$ in (4) as a single SDE

with

Let us note $\varrho_t(x,y)$ the density of $(X_t, W_t)$, and write the associated Fokker-Planck equation; here and after, $\nabla$ is a gradient and $\nabla \cdot$ is the divergence. For clarity in this section, subscripts denote partial differentiation with respect to the corresponding variables.

Now by the definition of $f_t$, we have $\nabla_{x,w}f_t(x,w) = - \nabla_x \nabla_x U_t(x) - \dot{U}_t(x)$, hence finally

The marginal integration

We already set $q_t(x) = \int_0^\infty w \varrho_t(x,w)dw$, so that $\mathbf{E}[W_t] = \int q_t(x)dx$ for example. Multiplying the last equation by $w$, integrating, and swapping integrals and derivatives, we get

We have $f_t \cdot \nabla_{x,w} \varrho_t = - \nabla_x U_t \cdot \nabla_x \varrho_t - \dot{U}_t w \nabla_w \varrho_t$, hence

An integration by parts shows that $\int_0^\infty w^2 \nabla_w \varrho_t(x,w)dw = -2 \int_0^\infty w \varrho_t(x,w)dw = -2q_t(x)$, provided a sufficient decay of $w \to w^2 \varrho(x,w)$ at infinity. We obtain the following equation:

where I removed the subscript for clarity. This equation is a Fokker-Planck equation (first terms) but with a birth-death term added at the end.

It turns out that the density $p_t(x) = e^{-U_t(x)}$ is a solution of this equation. Indeed, $\dot{p}_t = - \dot{U}_t p_t$ and $\nabla p_t = -\nabla U_t p_t$, so that it is easy to check that (16) is satisfied.

Uniqueness of solutions to parabolic PDEs

We're now left with checking that the solutions of

are essentially unique if they start with the same initial conditions, which is the case here. In the equation $B_t = \nabla U_t$ and $C_t=-\dot{U}_t$. Uniqueness for solutions of parabolic equations is well-known to PDE theorists and can be found in Evans' book, chapter 7. We provide a proof for completeness. Indeed, by linearity, it is sufficient to check that the only solution of (17) started at $u_0(x) = 0$ is identically zero. Set $\gamma(t):=\int |u_t(x)|^2dx$; we're going to bound its derivative using "energy estimates". We have

The first term is $-\int |\nabla u_t|^2$. The second term is equal to $+\int \nabla u_t \cdot u_t B_t$ by integration by parts. Young's inequality $a \cdot b \leqslant \lambda |a|^2 + |b|^2 / 4\lambda$, valid for any $\lambda>0$, yields

In particular with $\lambda=1$ we cancel the first term and we obtain

Now, if $B_t, C_t$ are bounded by some constant $c$, we get $\dot{\gamma}(t) \leqslant c \gamma(t)$. Grönwall's lemma yields $\gamma(t)\leqslant \gamma(0)e^{ct}=0$ and since $\gamma(t)$ is always positive, we get $\gamma(t)=0$ and $u_t=0$ for all $t$.

References

• Jarzynski's paper, written in 1996.