Malliavin Weight Sampling: A Practical Guide

Abstract

Malliavin weight sampling (MWS) is a stochastic calculus technique for computing the derivatives of averaged system properties with respect to parameters in stochastic simulations, without perturbing the system’s dynamics. It applies to systems in or out of equilibrium, in steady state or time-dependent situations, and has applications in the calculation of response coefficients, parameter sensitivities and Jacobian matrices for gradient-based parameter optimisation algorithms. The implementation of MWS has been described in the specific contexts of kinetic Monte Carlo and Brownian dynamics simulation algorithms. Here, we present a general theoretical framework for deriving the appropriate MWS update rule for any stochastic simulation algorithm. We also provide pedagogical information on its practical implementation.

1. Introduction

Malliavin weight sampling (MWS) is a method for computing derivatives of averaged system properties with respect to parameters in stochastic simulations [1,2]. The method has been used in quantitative financial modelling to obtain the “Greeks” (price sensitivities) [3], and as the Girsanov transform, in kinetic Monte Carlo simulations for systems biology [4]. Similar ideas have been used to study fluctuation-dissipation relations in supercooled liquids [5]. However, MWS appears to be relatively unknown in the fields of soft matter, chemical and biological physics, perhaps because the theory is relatively impenetrable for non-specialists, being couched in the language of abstract mathematics (e.g., martingales, Girsanov transform, Malliavin calculus, etc.); an exception in financial modelling is [6].

MWS works by introducing an auxiliary stochastic quantity, the Malliavin weight, for each parameter of interest. The Malliavin weights are updated alongside the system’s usual (unperturbed) dynamics, according to a set of rules. The derivative of any system function, A, with respect to a parameter of interest is then given by the average of the product of A with the relevant Malliavin weight; or in other words, by a weighted average of A, in which the weight function is given by the Malliavin weight. Importantly, MWS works for non-equilibrium situations, such as time-dependent processes or driven steady states. It thus complements existing methods based on equilibrium statistical mechanics, which are widely used in soft matter and chemical physics.

MWS has so far been discussed only in the context of specific simulation algorithms. In this paper, we present a pedagogical and generic approach to the construction of Malliavin weights, which can be applied to any stochastic simulation scheme. We further describe its practical implementation in some detail using as our example one dimensional Brownian motion in a force field.

2. The Construction of Malliavin Weights

The rules for the propagation of Malliavin weights have been derived for the kinetic Monte-Carlo algorithm [4,7], for the Metropolis Monte-Carlo scheme [5] and for both underdamped and overdamped Brownian dynamics [8]. Here we present a generic theoretical framework, which encompasses these algorithms and also allows extension to other stochastic simulation schemes.

We suppose that our system evolves in some state space, and a point in this state space is denoted as S. Here, we assume that the state space is continuous, but our approach can easily be translated to discrete or mixed discrete-continuous state spaces. Since the system is stochastic, its state at time t is described by a probability distribution, P (S). In each simulation step, the state of the system changes according to a propagator, W (S → S′), which gives the probability that the system moves from point S to point S′ during an application of the update algorithm. The propagator has the property that

$$P^{\prime }(S^{\prime })={\int }_{S}dS\hspace{0.17em}W(S\to S^{\prime })\hspace{0.17em}P(S)$$

(1)

where P′ (S) is the probability distribution after the update step has been applied and the integral is over the whole state space. We shall write this in a shorthand notation as

$$P^{\prime }=\int WP.$$

(2)

Integrating Equation (1) over S′, we see that the propagator must obey ∫_S′ W (S → S′) = 1. It is important to note, however, that we do not assume the detailed balance condition P_eq(S) W (S → S′) = P_eq(S′) W (S′ → S) (for some equilibrium P_eq(S)). Thus, our results apply to systems whose dynamical rules do not obey detailed balance (such as chemical models of gene regulatory networks [9]), as well as to systems out of steady state. We observe that the (finite) product

$$\text{𝕎}(S_1,\dots ,S_n)=W(S_1\to S_2)\times \cdots \times W(S_{n-1}\to S_n)$$

(3)

is proportional to the probability of occurrence of a trajectory of states, {S₁, . . . , S_n}, and can be interpreted as a trajectory weight.

Let us now consider the average of some quantity A(S) over the state space, in shorthand

$$\langle A\rangle =\int A\hspace{0.17em}P.$$

(4)

The quantity A might well be a complicated function of the state of the system: for example the extent of crystalline order in a particle-based simulation, or a combination of the concentrations of various chemical species in a simulation of a biochemical network. We suppose that we are interested in the sensitivity of 〈A〉 to variations in some parameter of the simulation, which we denote as λ. This might be one of the force field parameters (or the temperature) in a particle-based simulation or a rate constant in a kinetic Monte Carlo simulation. We are interested in computing ∂〈A〉/∂λ. This quantity can be written as

$$\frac{\partial \langle A\rangle }{\partial \lambda }=\int A\hspace{0.17em}P\hspace{0.17em}{Q}_{\lambda },$$

(5)

where

$$Q_{\lambda }=\frac{\partial \hspace{0.17em}\text{ln}\hspace{0.17em}P}{\partial \lambda }$$

(6)

(using the fact that ∂ ln P/∂λ = (1/P)∂P/∂λ).

Let us now suppose that we track in our simulation not only the physical state of the system, but also an auxiliary stochastic variable, which we term q_λ. At each simulation step q_λ is updated according to a rule that depends on the system state; this does not perturb the system’s dynamics, but merely acts as a “readout”. By tracking q_λ, we extend the state space, so that S becomes {S, q_λ}. We can then define the average 〈q_λ〉_S, which is an average of the value of q_λ in the extended state space, with the constraint that the original (physical) state space point is fixed at S (see further below).

Our aim is to define a set of rules for updating q_λ, such that 〈q_λ〉_S = Q_λ, i.e., such that the average of the auxiliary variable, for a particular state space point, measures the derivative of the probability distribution with respect to the parameter of interest, λ. If this is the case then, from Equation (5)

$$\frac{\partial \langle A\rangle }{\partial \lambda }=\langle A{q}_{\lambda }\rangle .$$

(7)

The auxiliary variable q_λ is the Malliavin weight corresponding to the parameter λ.

How do we go about finding the correct updating rule? If the Malliavin weight exists, we should be able to derive its updating rule from the system’s underlying stochastic equations of motion. We obtain an important clue from differentiating Equation (1) with respect to λ. Extending the shorthand notation, one finds

$$P^{\prime }{Q}_{\lambda }^{\prime }=\int WP\left(Q_{\lambda }+\frac{\partial \text{ln}W}{\partial \lambda }\right).$$

(8)

This strongly suggests that the rule for updating the Malliavin weight should be

$$q_{\lambda }^{\prime }=q_{\lambda }+\frac{\partial \hspace{0.17em}\text{ln}\hspace{0.17em}W}{\partial \lambda }.$$

(9)

In fact, this is correct. The proof is not difficult and, for the case of Brownian dynamics, can be found in the supplementary material for [8]. It involves averaging Equation (9) in the extended state space {S, q_λ}.

From a practical point of view, for each time step, we implement the following procedure:

• propagate the system from its current state S to a new state S′ using the algorithm that implements the stochastic equations of motion (Brownian, kinetic Monte-Carlo, etc.);

• with knowledge of S and S′, and the propagator W (S → S′), calculate the change in the Malliavin weight Δq_λ = ∂ ln W (S → S′)/∂λ;

• update the Malliavin weight according to q_λ → q′_λ = q_λ+ Δq_λ.

At the start of the simulation, the Malliavin weight is usually initialised to q_λ = 0.

Let us first suppose that our system is not in steady state. However rather the quantity 〈A〉 in which we are interested is changing in time and likewise ∂〈A(t)〉/∂λ is a time-dependent quantity. To compute ∂〈A(t)〉/∂λ we run N independent simulations, in each one tracking as a function of time A(t), q_λ(t) and the product A(t) q_λ(t). The quantities 〈A(t)〉 and ∂〈A(t)〉/∂λ are then given by

$$\langle A(t)\rangle \approx \frac{1}{N}\sum _{i=1}^N{A}_i(t),\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{\partial \langle A(t)\rangle }{\partial \lambda }\approx \frac{1}{N}\sum _{i=1}^N{A}_i(t)\hspace{0.17em}{q}_{\lambda ,i}(t),\hspace{0.17em}$$

(10)

where A_i(t) is the value of A(t) recorded in the ith simulation run (and likewise for q_λ,i(t)). Error estimates can be obtained from the variance among the replicate simulations.

If, instead, our system is in steady state, the procedure needs to be modified slightly. This is because the variance in the values of q (t) across replicate simulations increases linearly in time (this point is discussed further below). For long times, computation of ∂〈A〉/∂λ using Equation (10) therefore incurs a large statistical error. Fortunately, this problem can easily be solved by computing the correlation function

$$C(t,t^{\prime })=\langle A(t)[q_{\lambda }(t)-q_{\lambda }(t^{\prime })]\rangle .$$

(11)

In steady state, C(t, t′) = C(t − t′), with the property that C(Δt) → ∂A/∂λ as Δt → ∞. In a single simulation run, we simply measure q_λ(t) and A(t) at time intervals separated by Δt (which is typically multiple simulation steps). At each measurement, we compute A(t) [q_λ(t) − q_λ(t − Δt)]. We then average this latter quantity over the whole simulation run to obtain an estimate of ∂〈A〉/∂λ. For this estimate to be accurate, we require that Δt is long enough that C(Δt) has reached its plateau value; this typically means that Δt should be longer than the typical relaxation time of the system’s dynamics. The correlation function approach is discussed in more detail in [7,8].

Returning to a more theoretical perspective, it is interesting to note that the rule for updating the Malliavin weight, Equation (9), depends deterministically on S and S′. This implies that the value of the Malliavin weight at time t is completely determined by the trajectory of system states during the time interval 0 → t. In fact, it is easy to show that

$$q_{\lambda }=\frac{\partial \hspace{0.17em}\ln \hspace{0.17em}𝕎}{\partial \lambda }$$

(12)

where 𝕎 is the trajectory weight defined in Equation (3). Similar expressions are given in [5,7]. Thus, the Malliavin weight q_λ is not fixed by the state point S but by the entire trajectory of states that have led to state point S. Since many different trajectories can lead to S, many values of q_λ are possible for the same state point S. The average 〈q_λ(t)〉_S is actually the expectation value of the Malliavin weight, averaged over all trajectories that reach state point S at time t. This can be used to obtain an alternative proof that 〈q_λ〉_S = ∂ ln P/∂λ. Suppose we sample N trajectories, of which N_S end up at state point S (or a suitably defined vicinity thereof, in a continuous state space). We have P(S) = 〈N_S〉/N. Then, the Malliavin property implies ∂P/∂λ = 〈N_S q_λ〉/N, and hence ∂ ln P/∂λ = 〈N_S q_λ〉/〈N_S〉 = 〈q_λ〉_S.

3. Multiple Variables, Second Derivatives and the Algebra of Malliavin Weights

Up to now, we have assumed that the quantity A does not depend explicitly on the parameter λ. There may be cases, however, when A does have an explicit λ-dependence. In these cases, Equation (7) should be replaced by

$$\frac{\partial \langle A\rangle }{\partial \lambda }=\langle \frac{\partial \langle A\rangle }{\partial \lambda }\rangle +\langle A\hspace{0.17em}{q}_{\lambda }\rangle .$$

(13)

If we set A to be a constant in this, we immediately obtain the general result that 〈q_λ〉 = 0. Equation (13) reveals a kind of ‘algebra’ for Malliavin weights: we see that the operations of taking an expectation value and taking a derivative can be commuted, provided the Malliavin weight is introduced as the commutator.

We can also extend our analysis further to allow us to compute higher derivatives with respect to the parameters. These may be useful, for example, for increasing the efficiency of gradient-based parameter optimisation algorithms. Taking the derivative of Equation (13) with respect to a second parameter μ gives

$$\begin{array}{ll}\frac{{\partial }^2\langle A\rangle }{\partial \lambda \partial \mu }\hfill & =\frac{\partial }{\partial \mu }\langle \frac{\partial A}{\partial \lambda }\rangle +\frac{\partial \langle A\hspace{0.17em}{q}_{\lambda }\rangle }{\partial \mu }\hfill \\ \hfill & =\langle \frac{{\partial }^{2}A}{\partial \lambda \partial \mu }\rangle +\langle \frac{\partial A}{\partial \lambda }{q}_{\mu }\rangle +\langle A\frac{\partial q_{\lambda }}{\partial \mu }\rangle +\langle \frac{\partial A}{\partial \mu }{q}_{\lambda }\rangle +\langle A\hspace{0.17em}{q}_{\lambda }{q}_{\mu }\rangle \hfill \\ \hfill & =\langle A(q_{\lambda \mu }+q_{\lambda }{q}_{\mu })\rangle +\langle \frac{\partial A}{\partial \lambda }{q}_{\mu }\rangle +\langle \frac{\partial A}{\partial \mu }{q}_{\lambda }\rangle +\langle \frac{{\partial }^{2}A}{\partial \lambda \partial \mu }\rangle ,\hfill \end{array}$$

(14)

where in the second line we iterate the commutation relation, and in the third line we collect like terms and introduce

$$q_{\lambda \mu }=\frac{\partial q_{\lambda }}{\partial \mu }.$$

(15)

In the case where A is independent of the parameters, this result simplifies to

$$\frac{{\partial }^2\langle A\rangle }{\partial \lambda \partial \mu }=\langle A(q_{\lambda \mu }+q_{\lambda }{q}_{\mu })\rangle .$$

(16)

The quantity q_λμ here is a new, second order, Malliavin weight which from Equations (12) and (15) satisfies

$$q_{\lambda \mu }=\frac{{\partial }^2\hspace{0.17em}\ln \hspace{0.17em}𝕎}{\partial \lambda \partial \mu }.$$

(17)

To compute second derivatives with respect to the parameters, we should therefore track these second order Malliavin weights in our simulation, updating them alongside the existing Malliavin weights by the rule

$$q_{\lambda \mu }^{\prime }=q_{\lambda \mu }+\frac{{\partial }^2\hspace{0.17em}\ln \hspace{0.17em}W(S\to S^{\prime })}{\partial \lambda \partial \mu }.$$

(18)

Setting A as a constant in Equation (16), we also obtain the interesting result that 〈q_λμ〉 = −〈q_λq_μ〉.

Steady state problems can be approached by extending the correlation function method to second order weights. Define, cf. Equation (11),

$$C(t,\hspace{0.17em}{t}^{\prime })=\langle A(t)\left\{[q_{\lambda \mu }(t)+q_{\lambda }(t)q_{\mu }(t)]-[q_{\lambda \mu }(t^{\prime })+q_{\lambda }(t^{\prime })q_{\mu }(t^{\prime })]\right\}\rangle .$$

(19)

As in the first order case, in steady state we expect C(t, t′) = C(t − t′) with the property that C(Δt) → ∂²〈A〉/∂λ∂μ as Δt → ∞.

4. One-Dimensional Brownian Motion in a Force Field

We now demonstrate this machinery by way of a practical but very simple example, namely one-dimensional (overdamped) Brownian motion in a force field. In this case, the state space is specified simply by the particle position x which evolves according to the Langevin equation

$$\frac{dx}{dt}=f(x)+\eta$$

(20)

where f(x) is the force field and η is Gaussian white noise of amplitude 2T, where T is temperature. Without loss of generality, we have chosen units so that there is no prefactor multiplying the force field. We discretise the Langevin equation to the following updating rule:

$$x^{\prime }=x+f(x)\delta t+\xi ,$$

(21)

where δt is the time step and ξ is a Gaussian random variate with zero mean and variance 2T δt. Corresponding to this updating rule is an explicit expression for the propagator,

$$W(x\to x^{\prime })=\frac{1}{\sqrt{4\pi T\delta t}}\exp \left(-\frac{{\left(x^{\prime }-x-f(x)\delta t\right)}^2}{4T\hspace{0.17em}\delta t}\right).$$

(22)

This follows from the statistical distribution of ξ. Let us suppose that the parameter of interest λ enters into the force field (the temperature T could also be chosen as a parameter). Making this assumption

$$\frac{\partial \ln W(x\to x^{\prime })}{\partial \lambda }=\frac{(x^{\prime }-x-f\delta t)}{2T}\frac{\partial f}{\partial \lambda }.$$

(23)

We can simplify this result by noting that from Equation (21), x′ − x − f δt = ξ. Making use of this, the final updating rule for the Malliavin weight is

$$q_{\lambda }^{\prime }=q_{\lambda }+\frac{\xi }{2T}\frac{\partial f}{\partial \lambda }$$

(24)

where ξ is the exact same value that was used for updating the position in Equation (21). Because the value of ξ is the same for the updates of position and of q_λ, the change in q_λ is completely determined by the end points, x and x′. The derivative ∂f/∂λ should be evaluated at x since that is the position at which the force is computed in Equation (21). Since ξ in Equation (21) is a random variate uncorrelated with x, averaging Equation (24) shows that 〈q′_λ 〉 = 〈q_λ 〉. As the initial condition is q_λ = 0, this means that 〈q_λ〉 = 0, as predicted in the previous section. Equation (24) is essentially the same as that derived in [8].

If we differentiate Equation (23) with respect to a second parameter μ we get

$$\frac{{\partial }^2\ln W(x-x^{\prime })}{\partial \lambda \partial \mu }=\frac{(x^{\prime }-x-f\hspace{0.17em}\delta t)}{2T}\frac{{\partial }^{2}f}{\partial \lambda \partial \mu }-\frac{\delta t}{2T}\frac{\partial f}{\partial \lambda }\frac{\partial f}{\partial \mu }.$$

(25)

Hence, the updating rule for the second order Malliavin weight can be written as

$$q_{\lambda \mu }^{\prime }=q_{\lambda \mu }+\frac{\xi }{2T}\frac{{\partial }^{2}f}{\partial \lambda \partial \mu }-\frac{\delta t}{2T}\frac{\partial f}{\partial \lambda }\frac{\partial f}{\partial \mu },$$

(26)

where again ξ is the exact same value as that used for updating the position in Equation (21). If we average Equation (26) over replicate simulation runs, we find 〈q′_λμ〉 = 〈q_λμ〉−(δt/2T)(∂f/∂λ)(∂f/∂μ). Hence the mean value 〈q_λμ〉 drifts in time, unlike 〈q_λ〉 or 〈q_μ〉. However, one can show that the mean value of the sum 〈(q_λμ + q_λq_μ)〉 is constant in time and equal to zero as long as initially q_λ = q_μ = 0.

Now let us consider the simplest case of a particle in a linear force field f = −κx + h (also discussed in [8]). This corresponds to a harmonic trap with the potential $U=\frac{1}{2}\kappa x^2-hx$. We let the particle start from x₀ at t = 0 and track its time-dependent relaxation to the steady state. We shall set T = 1 for simplicity. The Langevin equation can be solved exactly for this case, and the mean position evolves according to

$$\langle x(t)\rangle =x_0{e}^{-\kappa t}+\frac{h}{\kappa }(1-e^{-\kappa t}).$$

(27)

We suppose that we are interested in derivatives with respect to both h and κ, for a “baseline” parameter set in which κ is finite but h = 0. Taking derivatives of Equation (27) and setting h = 0, we find

$$\frac{\partial \langle x(t)\rangle }{\partial h}=\frac{1-e^{-\kappa t}}{\kappa },\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{\partial \langle x\rangle (t)}{\partial \kappa }=-x_{0}t{e}^{-\kappa t},\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{{\partial }^2\langle x(t)\rangle }{\partial h\partial \kappa }=\frac{t{e}^{-\kappa t}}{\kappa }-\frac{1-e^{-\kappa t}}{{\kappa }^2}\hspace{0.17em}.$$

(28)

We now show how to compute these derivatives using Malliavin weight sampling. Applying the definitions in Equations (24) and (26), the Malliavin weight increments are

$$q_h^{\prime }=q_h+\frac{\xi }{2},\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}{q}_{\kappa }^{\prime }=q_{\kappa }-\frac{x\xi }{2},\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}{q}_{h\kappa }^{\prime }=q_{h\kappa }+\frac{x\delta t}{2},$$

(29)

and the position update itself is

$$x^{\prime }=x-\kappa x\hspace{0.17em}\delta t+\xi .$$

(30)

We track these Malliavin weights in our simulation and use them to calculate derivatives according to

$$\frac{\partial \langle x(t)\rangle }{\partial h}=\langle x(t)q_h(t)\rangle ,\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{\partial \langle x(t)\rangle }{\partial \kappa }=\langle x(t)q_{\kappa }(t)\rangle ,\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{{\partial }^2\langle x(t)\rangle }{\partial h\partial \kappa }=\langle x(t)\left(q_{h\kappa }(t)+q_h(t)q_{\kappa }(t)\right)\rangle .$$

(31)

Equations (29)–(31) have been coded up as a MATLAB script, described in Section 5. A typical result generated by running this script is shown in Figure 1. Equations (29) and (30) are iterated with δt = 0.01 up to t = 5, for a trap strength κ = 2 and initial position x₀ = 1. The weighted averages in Equation (31) are evaluated as a function of time for N = 10⁵ samples as in Equation (10). These results are shown as the solid lines in Figure 1. The dashed lines are theoretical predictions for the time dependent derivatives from Equation (28). As can be seen, the agreement between the time-dependent derivatives and the Malliavin weight averages is very good.

As discussed briefly above, in this procedure the sampling error in the computation of ∂〈A(t)〉/∂λ is expected to grow with time. Figure 2 shows the mean square Malliavin weight as a function of time for the same problem. For the first order weights q_h and q_κ the growth rate is typically linear in time. Indeed, from Equation (29), one can prove that in the limit δt → 0 (see Section 5)

$$\frac{d\langle q_h^2\rangle }{dt}=\frac{1}{2},\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\hspace{0.17em}\frac{d\langle q_{\kappa }^2\rangle }{dt}=\frac{\langle x^2\rangle }{2}.$$

(32)

Thus q_h behaves exactly as a random walk, as should be obvious from the updating rule. The other weight q_λ also ultimately behaves as a random walk since 〈x²〉 = 1/κ in steady state (from equipartition). Figure 2 also shows that the second order weight q_hκ grows superdiffusively; one can show that eventually 〈(q_hκ + q_hq_κ)²〉 ∼ t², although the transient behaviour is complicated. Full expressions are given in Section 5. This suggests that computation of second order derivatives is likely to suffer more severely from statistical sampling problems than the computation of first order derivatives.

5. Conclusions

In this paper, we have provided an outline of the generic use of Malliavin weights for sampling derivatives in stochastic simulations, with an emphasis on practical aspects. The usefulness of MWS for a particular simulation scheme hinges on the simplicity or otherwise of constructing the propagator W (S → S′) which fixes the updating rule for the Malliavin weights according to Equation (9). The propagator is determined by the algorithm used to implement the stochastic equations of motion; MWS may be easier to implement for some algorithms than for others. We note however that there is often some freedom of choice about the algorithm, such as the choice of a stochastic thermostat in molecular dynamics, or the order in which update steps are implemented. In these cases, a suitable choice may simplify the construction of the propagator and facilitate the use of Malliavin weights.