Title: Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transition Path Sampling

Protein Folding

Schrödinger Bridge
Abstract: Rare event sampling in dynamical systems is a fundamental problem arising in the natural sciences, which poses significant computational challenges due to an exponentially large space of trajectories. For settings where the dynamical system of interest follows a Brownian motion with known drift, the question of conditioning the process to reach a given endpoint or desired rare event is definitively answered by Doob's $h$-transform. However, the naive estimation of this transform is infeasible, as it requires simulating sufficiently many forward trajectories to estimate rare event probabilities. In this work, we propose a variational formulation of Doob's $h$-transform as an optimization problem over trajectories between a given initial point and the desired ending point. To solve this optimization, we propose a simulation-free training objective with a model parameterization that imposes the desired boundary conditions by design. Our approach significantly reduces the search space over trajectories and avoids expensive trajectory simulation and inefficient importance sampling estimators which are required in existing methods. We demonstrate the ability of our method to find feasible transition paths on real-world molecular simulation and protein folding tasks.

Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling

Abstract

Rare event sampling in dynamical systems is a fundamental problem arising in the natural sciences, which poses significant computational challenges due to an exponentially large space of trajectories. For settings where the dynamical system of interest follows a Brownian motion with known drift, the question of conditioning the process to reach a given endpoint or desired rare event is definitively answered by Doob's h -transform. However, the naive estimation of this transform is infeasible, as it requires simulating sufficiently many forward trajectories to estimate rare event probabilities. In this work, we propose a variational formulation of Doob's h -transform as an optimization problem over trajectories between a given initial point and the desired ending point. To solve this optimization, we propose a simulation-free training objective with a model parameterization that imposes the desired boundary conditions by design. Our approach significantly reduces the search space over trajectories and avoids expensive trajectory simulation and inefficient importance sampling estimators which are required in existing methods. We demonstrate the ability of our method to find feasible transition paths on real-world molecular simulation and protein folding tasks.

1 Introduction

Conditioning a stochastic process to obey a particular endpoint distribution, satisfy desired terminal conditions, or observe a rare event is a problem with a long history (Schrödinger, 1932; Doob, 1957) and wide-ranging applications from generative modeling (De Bortoli et al., 2021; Chen et al., 2021a; Liu et al., 2022, 2023c; Somnath et al., 2023) to molecular simulation (Anderson, 2007; Wu et al., 2022; Plainer et al., 2023; Holdijk et al., 2023), drug discovery (Kirmizialtin et al., 2012, 2015; Dickson, 2018), and materials science (Xi et al., 2013; Selli et al., 2016; Sharma et al., 2016).

Transition Path Sampling. In this work, we take a particular interest in the problem of transition path sampling (TPS) in computational chemistry (Dellago et al., 2002; Weinan and Vanden-Eijnden, 2010), which attempts to describe how molecules transition between local energy minima or metastable states under random fluctuations or the influence of external forces. Understanding such transitions has numerous applications for combustion, catalysis, battery, material design, and protein folding (Zeng et al., 2020; Klucznik et al., 2024; Blau et al., 2021; Noé et al., 2009; Escobedo et al., 2009).

While the TPS problem is often framed as finding the 'most probable path' transitioning between states (Dürr and Bach, 1978; Vanden-Eijnden and Heymann, 2008), we build upon connections between TPS and Doob's h -transform (Das et al., 2019, 2021, 2022; Koehl and Orland, 2022; Singh and Limmer, 2023) and seek to match the full posterior distribution over conditioned processes.

Doob's h -Transform. For Brownian motion diffusion processes, conditioning is known to be achieved by Doob's h -transform (Doob, 1957; Särkkä and Solin, 2019). However, solving this problem amounts to estimating rare event probabilities or matching a complex target distribution. Approaches which involve simulation of trajectories to construct Monte Carlo expectations or importance sampling estimators (Papaspiliopoulos and Roberts, 2012; Schauer et al., 2017; Yan et al., 2022; Holdijk et al., 2023) can be extremely inefficient if the target event is rare or endpoint distribution is difficult to match. Recent methods based on score matching (Heng et al., 2021) or nonlinear Feynman-Kac formula (Chopin et al., 2023) still require simulation during optimization.

Variational Formulation of Doob's h -Transform. In this work, we propose a variational formulation of Doob's h -transform as the solution to an optimization on the space of paths of probability distributions. We focus on solving for the Doob transform conditioning on a particular terminal point, which is natural in the TPS setting (see Fig. 1). Taking inspiration from recent bridge matching methods (Peluchetti, 2021, 2023; Liu et al., 2022; Lipman et al., 2022; Shi et al., 2023; Liu et al., 2023a), we propose a parameterization with the following attractive features.

Every Sample Matters. In contrast to most existing approaches, our training method is simulation-free , thereby avoiding computationally wasteful simulation methods to estimate rare-event probabilities and inefficient importance or rejection sampling. We thus refer to our approach as being sample-efficient .

Optimization over Sampling. We propose an expressive variational family of approximations to the conditioned process, which are tractable to sample and can be optimized using neural networks with end-to-end backpropagation.

Problem-Informed Parameterization. Our parameterization enforces the boundary conditions by design , thereby reducing the search space for optimization and efficiently making use of the conditioning information.

We begin by linking the problem of transition path sampling to the Doob's h -transform and recalling background results in Sec. 2. We present our variational formulation in Sec. 3.1 and detail our optimization algorithm throughout Sec. 3.2. We demonstrate the ability of our approach to achieve comparable performance to Markov Chain Monte Carlo (MCMC) methods with notably improved efficiency on synthetic, and real-world molecular simulation tasks in Sec. 5.

2 Background

2.1 Transition Path Sampling

Consider a forward or reference stochastic process with states x t and the density of transition probability ρ t + dt ( y | x t = x ) := ρ ( x t + dt = y | x t = x ) . Starting from an initial point x 0 = A , the probability density of a discrete-time path is given as

The problem of rare event sampling aims to condition this reference stochastic process on some event at time T , for example, that the final state belongs to a particular set x T ∈ B . We are interested in

sampling from the entire transition path , namely the posterior distribution over intermediate states

Moving to continuous time, we focus on the transition path sampling problem in the case where the reference process is given by a Brownian motion. In particular, we are motivated by applications in computational chemistry (Dellago et al., 2002; Weinan and Vanden-Eijnden, 2010), where the reference process is given by molecular dynamics following either overdamped Langevin dynamics,

or the second-order Langevin dynamics with spatial coordinates ¯ x t and velocities ¯ v t ,

for a potential energy function U , where W t denotes the Wiener process. Note that k B T is the Boltzman constant times temperature, M is the mass matrix, and γ is the friction coefficient.

2.2 Doob's h -transform

Doob's h -transform addresses the question of conditioning a reference Brownian motion to satisfy a terminal condition such as x T ∈ B , thereby providing an avenue to solve the transition path sampling problem described above. Without loss of generality, and to provide a unified treatment of the dynamics in (3)-(4), we consider the forward or reference stochastic differential equation (SDE),

with drift vector field b t : R N → R N and diffusion coefficient matrix Ξ t ∈ R N × N such that G t := 1 2 Ξ t Ξ T t is positive definite. 2 We denote the induced path measure as P ref 0: T ∈ P ( C ([0 , T ] → R N )) , i.e. a measure over continuous functions from time to R N .

Remarkably, Doob's h -transform (Doob, 1957; Särkkä and Solin, 2019, Sec. 7.5) shows that conditioning the reference process (5) on x T ∈ B results in another Brownian motion process.

Proposition 1. [Jamison (1975, Thm. 2)] Let h B ( x, t ) := ρ T ( x T ∈ B| x t = x ) denote the conditional transition probability of the reference process in (5) . Then,

where we use x t | T to denote a conditional process. The SDE in (6) is associated with the following transition probabilities for s < t < T ,

Note that all of our subsequent results hold for the case when B is a point-mass, with the only change being that the h -function becomes a density, h B ( x, t ) = ρ T ( B | x t = x ) .

See App. A.1 for proof, and note that (7) is simply an application of Bayes rule ρ t ( y | x s = x, x T ∈ B ) = ρ T ( x T ∈ B| x t = y ) ρ t ( y | x s = x ) /ρ T ( x T ∈ B| x s = x ) with the unconditioned or reference transition probability as the prior. Furthermore, the conditioned transition probabilities in (7) allow us to directly construct the transition path (2). Using Bayes rule, we have

after telescoping cancellation of h -functions and rewriting the denominator in (2) as h B ( A, 0) Thus, we can solve the TPS problem by exactly solving for the h -function and simulating the SDE in (6).

Finally, the h -process and temporal marginals ρ t ( x | x 0 = A,x T ∈ B ) of the conditioned process satisfy the following forward and backward Kolmogorov equations, which will be useful in deriving our variational objectives in the next section. Note, we use ⟨∇ x , · ⟩ = div ( · ) for the divergence operator, and we use ρ t | 0 ,T to indicate the dependence on both x 0 = A (via the initial condition of (8a)) and x T ∈ B (via the h -transform h B ). See App. A.1 for the proof.

Proposition 2. The following PDEs are obeyed by (a) the density of the conditioned process ρ t | 0 ,T ( x ) := ρ t ( x | x 0 = A,x T ∈ B ) and (b) the h -function h B ( x, t ) ,

Reparameterizing (8b) in terms of s B ( x, t ) := log h B ( x, t ) , we can also write

3 Method

We first present a novel variational objective whose minimum corresponds to the Doob h -transform in Sec. 3.1, and then propose an efficient parameterization to solve for the h -transform in Sec. 3.2.

3.1 Doob's Lagrangian

Consider reference dynamics given in the form of either (3) or (4), with known drift b t or energy U . We will restrict our attention to conditioning on a terminal rare event of reaching a given endpoint x T = B , along with an initial point x 0 = A . We approach solving for Doob's h -transform via a least action principle where, in the following theorem, we define a Lagrangian action whose minimization yields the optimal q ∗ t | 0 ,T ( x ) = ρ t | 0 ,T ( x ) and v ∗ t | 0 ,T ( x ) = ∇ x log h B ( x, t ) from Prop. 1 and 2.

Theorem 1. The following Lagrangian action functional has a unique solution which matches the Doob h -transform in Prop. 2,

The optimal q ∗ t | 0 ,T ( x ) obeys (8a) , and v ∗ t | 0 ,T ( x ) = ∇ x log h B ( x, t ) = ∇ x s B ( x, t ) obeys (8b) -(8c) .

This objective will form the basis for our computational approach, with proof of Thm. 1 deferred to App. A.2. We proceed briefly to contextualize our variational objective and highlight several optimization challenges which will be solved by our proposed parameterization in Sec. 3.2.

Unconstrained Dual Objective. Introducing Lagrange multipliers to enforce the constraints in (9b)-(9c) and eliminating v t | 0 ,T , we obtain an alternative, unconstrained version of (9a).

Corollary 1. The Lagrangian objective in Thm. 1 which solves Doob's h -transform is equivalent to

if q t | 0 ,T satisfies (9c) . Note v t | 0 ,T ( x ) = ∇ x s B ( x, t ) , with s ∗ B ( x, t ) = log h B ( x, t ) at optimality. 3

This objective is similar to the objectives optimized by Action Matching methods (Neklyudov et al., 2023, 2024). Notably, the objective in Cor. 1 is expressed directly in terms of the (log) of the h -function for fixed conditioning information x T = B . We also note that the Hamilton Jacobi-style quantity, whose expectation appears in the final term, is zero at optimality in (8c) of Prop. 2.

Path Measure Perspective. We next relate our variational objective in Thm. 1 to a KL divergence optimization over path measures. Let P ref 0: T denote the law of the reference SDE in (5) with fixed P ref 0 = δ ( x 0 -A ) . Let Q v 0: T denote the law of a controlled process similar to (6), but with a variational v t | 0 ,T in place of ∇ x log h B ,

Note that the density q t | 0 ,T of Q v 0: T evolve according to the Fokker-Planck equation in (9b) (Särkkä and Solin, 2019, Sec. 5.2) . Using the Girsanov Theorem, the objective in (9a) can then be viewed as a KL divergence minimization over path measures Q v 0: T which satisfy the boundary constraints.

Corollary 2. The following Schrödinger Bridge (SB) problem

yields the path measure P ∗ 0: T associated with the SDE in (6) as its unique minimizing argument. The temporal marginals of P ∗ 0: T are equal to those which optimize the Lagrangian objective in Thm. 1.

Our Lagrangian action minimization thus corresponds to the solution of an SB problem (Schrödinger, 1932; Léonard, 2014) with Dirac delta functions as the endpoint measures. Our objective in (9a) particularly resembles optimal control formulations of SB (Chen et al., 2016, 2021b, Prob. 4.4, 5.3). While it is well-known that the Doob h -transform (and large deviation theory more generally) plays a role in the solution to SB problems (Jamison, 1975; Léonard, 2014), our interest in the transition path sampling problem leads to specific computational decisions below. See Sec. 4 for further discussion.

Challenges of Optimizing (9a). We highlight several distinctive features of our problem which inform the development of new computational methods in Sec. 3.2.

First, we perform optimization over the first argument of the KL divergence in (11), indicating that we need to be able to efficiently sample from the conditioned process in (10) or q t | 0 ,T in (9). This appears challenging due to the nonlinearity of both the reference and variational drifts, b t and v t | 0 ,T .

For a given q t | 0 ,T , it can be difficult to solve for v t | 0 ,T which satisfies the FokkerPlanck equation in (9b) or ∇ s which solves the inner optimization in Cor. 1.

̸

Finally, we would like to strictly enforce the boundary constraints on q t | 0 ,T or Q v 0: T to avoid simulating or wasting computation on trajectories for which x T = B .

In fact, our parameterization of q t | 0 ,T in Sec. 3.2 will completely avoid simulation of the SDE in (10) during training (Challenge 1), provide analytic solutions for v t | 0 ,T satisfying (9b) with a given q t | 0 ,T (Challenge 2), and exactly enforce the boundary constraints (Challenge 3).

3.2 Computational Approach

We now propose a family of Gaussian (mixture) path parameterizations q t | 0 ,T which overcome the computational challenges posed in the previous section, while still maintaining expressivity. We present all aspects of our proposed method in the context of the first-order dynamics (3) in Sec. 3.2.1, before presenting extensions to mixture paths and the second-order setting (4) in Sec. 3.2.2-3.2.3.

3.2.1 First-Order Dynamics and General Approach

Tractable Drift v t | 0 ,T for Variational Doob Objective. We begin by considering a modification of the Fokker-Planck constraint in (9b), with all drift terms absorbed into a single vector field u t | 0 ,T ,

For arbitrary q t | 0 ,T , solving for any u t | 0 ,T ( x ) satisfying (12) can be a difficult optimization problem, whose solution is not unique without some cost-minimizing assumption (Neklyudov et al., 2023).

To sidestep this optimization, and address Challenge 2, we restrict attention to variational families of q t | 0 ,T ∈ Q where it is analytically tractable to calculate a vector field u ( q,θ ) t | 0 ,T which satisfies (12). We first consider the family of Gaussian paths Q G , in similar fashion to (conditional) flow matching methods (Lipman et al., 2022; Tong et al., 2023; Liu et al., 2023a), with proof in App. B.

Algorithm 1: Training (Single Gaussian)

Algorithm 2: Sampling Trajectories

Algorithms for training with a single Gaussian path (Alg. 1) and sampling or generating transition paths at test time (Alg. 2). Note that we sample from the marginals q t | 0 ,T during training, but generate paths by simulating the SDE (10).

Proposition 3. For the family of endpoint-conditioned marginals q t | 0 ,T ( x ) = N ( x | µ t | 0 ,T , Σ t | 0 ,T ) ,

satisfies the Fokker-Planck equation (12) for q t | 0 ,T and diffusion coefficients G t = 1 2 Ξ t Ξ T t .

Given u ( q,θ ) t | 0 ,T corresponding to q t | 0 ,T , we can simply solve for the v t | 0 ,T satisfying the Fokker-Planck equation in (9b) in our variational Doob objective (Thm. 1). Since G t was assumed to be invertible and the base drift b t is known, we have

We may now evaluate terms involving v t | 0 ,T in our Lagrangian objective in (9) using (14) directly, without spending effort to solve an inner minimization over v t | 0 ,T (thus addressing Challenge 2).

Optimization over q t | 0 ,T satisfying Boundary Constraints. Given the ability to evaluate v ( q,θ ) t | 0 ,T for a given q t | 0 ,T ∈ Q G as above, our variational Doob objective in (9a) reduces to a single optimization over the marginals q t | 0 ,T of a conditioned process which satisfies the boundary conditions (9c).

We consider parameterizing the mean µ t | 0 ,T and covariance Σ t | 0 ,T of our Gaussian path q t | 0 ,T using a neural network. For simplicity, we consider a diagonal parameterization Σ t | 0 ,T = diag ( { σ 2 t | 0 ,T,d } D d =1 ) . We parameterize a neural network NNET θ : [0 , T ] × R D × R D → R D × R D which inputs time t and boundary conditions x 0 = A,x T = B , and outputs vectors of mean perturbations and per-dimension variances. Finally, using index notation to separate the output, we construct

Crucially, our Gaussian parameterization addresses Challenge 1, in that we can easily draw samples x t | 0 ,T ∼ q t | 0 ,T from our variational conditioned process (9b) without simulating the corresponding SDE with nonlinear drift (10). Further, the coefficients in (15b) and (15c) ensure that, as t → 0 or t → T , our parameterization satisfies the (smoothed) boundary conditions by design (Challenge 3). Although we add σ 2 min to ensure invertibilty of Σ t | 0 ,T (see (13)), we preserve q 0 ( x 0 ) = N ( x 0 | A,σ 2 min I D ) ≈ δ ( x 0 -A ) and q T ( x T ) = N ( x T | B,σ 2 min I D ) ≈ δ ( x T -B ) .

Reparameterization Gradients. Having shown that our parameterization satisfies the constraints (9b)-(9c) by design, we can finally optimize our variational Doob objective with respect to q t | 0 ,T ∈ Q G using the reparameterization trick (Kingma and Welling, 2013; Rezende et al., 2014). In particular, for the expectation at each t in (9a), we rewrite

where x = g ( t, ϵ ; θ ) is the mapping in (15) and v ( q,θ ) t | 0 ,T depends on θ via µ ( θ ) t | 0 ,T , Σ ( θ ) t | 0 ,T in (13)-(14).

Full Training Algorithm. In practice, we sample a batch of times { t i } i M =1 uniformly from the interval t ∈ [0 , T ] . For each time point, we approximate the gradient using a single-sample estimate of the expectation above (or (9)), which yields a simulation-free training procedure. The full training algorithm is outlined in Alg. 1.

Sampling of Trajectories. While we sample directly from q ( θ ) t | 0 ,T during training, we can sample full trajectories which obey this sequence of marginals at test time (Alg. 2). In particular, we simulate SDE trajectories with drift u ( q,θ ) t | 0 ,T ( x ) and diffusion coefficient G t using an appropriate solver. Note that this generation scheme sidesteps computationally expensive evaluation of the force field or base drift b t ( x t ) . We visualize example sampling trajectories in Fig. 2.

3.2.2 Second-Order Dynamics

To handle the case of the second-order dynamics in (4), we can adapt our recipe from the previous section with minimal modifications by extending the state space x ∈ R D to include velocities ¯ v , with x = (¯ x, ¯ v ) ∈ R 2 D . However, note that the dynamics in (4) are no longer stochastic in the spatial coordinates ¯ x . To ensure invertibility of G t and existence of the h -transform, we add a small nonzero diffusion coefficient in the coordinate space ¯ x , so that the reference process in Eq. (5) is given by

All steps in our algorithm proceed in similar fashion to Sec. 3.2.1. We now parameterize q t | 0 ,T (¯ x, ¯ v ) using NNET θ : [0 , T ] × R 2 D × R 2 D → R 2 D × R 2 D , which outputs mean perturbations and perdimension variances to calculate µ ¯ x t | 0 ,T , µ ¯ v t | 0 ,T and Σ ¯ x t | 0 ,T , Σ ¯ v t | 0 ,T and sample (¯ x, ¯ v ) , as in (15). Note that we parameterize Σ ¯ x t | 0 ,T , Σ ¯ v t | 0 ,T separately, matching the block diagonal form of (16). We calculate v ( q ) t | 0 ,T (¯ x, ¯ v ) := [ v ¯ x ( q ) t | 0 ,T , v ¯ v ( q ) t | 0 ,T ] from u ( q ) t | 0 ,T (¯ x, ¯ v ) = [ u ¯ x ( q ) t | 0 ,T , u ¯ v ( q ) t | 0 ,T ] as in (13)-(14), with G -1 t = ( 1 2 Ξ t Ξ T t ) -1 given by (16). The Lagrangian objective in (9) minimizes the norm of the concatenated vector v ( q ) t | 0 ,T (¯ x, ¯ v ) , which depends on the reference drift b t (¯ x, ¯ v ) in (16).

3.2.3 Gaussian Mixture Paths

Note that the true Doob h -transform may not yield marginal distributions which are unimodal Gaussians as in the previous section. To increase the expressivity of our variational family of conditioned processes, we now extend our parameterization to mixtures of Gaussians, q t | 0 ,T ∈ Q K MoG . Given a set of K mixture weights w := { w k } K k =1 and component Gaussian paths { q k t | 0 ,T } K k =1 , the following identity allows us to obtain the drift u ( q,θ ) of the corresponding mixture distribution q t | 0 ,T t | 0 ,T .

Proposition 4. Given a set of processes q k t | 0 ,T ( x ) and mixtures weights w k , the vector field satisfying the Fokker-Planck equation in (12) for the mixture q t | 0 ,T ( x ) = ∑ k w k q k t | 0 ,T ( x ) is given by

where u ( q,k ) t | 0 ,T ( x ) satisfies the Fokker-Planck equation in (12) for q k t | 0 ,T ( x ) . This identity holds for both first-order dynamics in spatial coordinates only or second-order dynamics in x = (¯ x, ¯ v ) .

Finally, we can calculate v ( q,θ ) t | 0 ,T ( x ) by comparing u ( q,θ ) t | 0 ,T ( x ) for the mixture of Gaussian path q t | 0 ,T ∈ Q K MoG to the reference drift b t ( x ) as in (14), and proceed to minimize its norm as in (9). In practice, we use Gumbel softmax reparamerization gradients (Maddison et al., 2016; Jang et al., 2017) to optimize the mixture weights { w k } K k =1 alongside the neural network parameters { θ k } K k =1 for each Gaussian component { µ ( θ k ) t | 0 ,T , Σ ( θ k ) t | 0 ,T } K k =1 and either first- or second-order dynamics.

4 Related Work

(Aligned) Schrödinger Bridge Matching Methods. Many existing 'bridge matching' approaches (Shi et al., 2023; Peluchetti, 2021, 2023; Liu et al., 2022; Lipman et al., 2022; Liu et al., 2023b) for SB and generative modeling rely on convenient properties of Brownian bridges and would require calculating h -transforms to simulate bridges for general reference processes. Our conditional

Gaussian path parameterization is similar to Liu et al. (2023a); Neklyudov et al. (2024), where analytic bridges are not available for SB problems with nonlinear reference drift or general costs.

Somnath et al. (2023); Liu et al. (2023b) attempt to solve the SB problem given access to aligned data x 0 , x T ∼ q data 0 ,T assumed to be drawn from an optimal coupling. While the method in Somnath et al. (2023) involves approximating an h -transform, their goal is to obtain an unconditioned vector field v t to simulate a Markov process. However, De Bortoli et al. (2023) use Doob's h -transform to argue the learned Markov process will not preserve the empirical coupling unless q data 0 ,T is the optimal coupling for the SB problem, and show that an 'augmented' v 0 ,t which conditions on x 0 can correct this issue.

After training on a dataset of x 0 , x T ∼ q data 0 ,T pairs using our method, we could consider using an (augmented) bridge matching objective (Shi et al., 2023; De Bortoli et al., 2023) to distill our learned v ( q ) t | 0 ,T into a vector field v t or v 0 ,t which does not condition on the endpoint. Our use of a Gaussian path parameterization with samples from a fixed endpoint coupling and no Markovization step corresponds to a simplified version of the conditional optimal control step in Liu et al. (2023a).

Transition Path Sampling. We refer to the surveys of Dellago et al. (2002); Weinan and VandenEijnden (2010); Bolhuis and Swenson (2021) for an overview of the TPS problem. Least action principles for TPS have a long history, building upon the Freidlin-Wentzell (Freidlin and Wentzell, 1998) and Onsager-Machlup (Onsager and Machlup, 1953; Dürr and Bach, 1978) Lagrangian functionals in the zero-noise limit and finite-noise cases. In particular, the Onsager-Machlup functional relates maximum a posteriori estimators or 'most probable (conditioned) paths' to the minimizers of an action functional similar to Thm. 1, where example algorithms include (Vanden-Eijnden and Heymann, 2008; Sheppard et al., 2008). By contrast, our approach targets the entire posterior over transition paths using an expressive variational family. While Lu et al. (2017) provide analysis for the Gaussian family, we draw connections with Doob's h -transform and extend to mixtures of Gaussians.

Shooting methods are among the most popular for sampling the posterior of transition paths. From a path that satisfies the boundary conditions (obtained, e.g., using high-temperature simulations), shooting picks points and directions to propose alterations, then simulates new trajectories and accepts or rejects using Metropolis-Hastings (MH) (Juraszek and Bolhuis, 2008; Borrero and Dellago, 2016; Jung et al., 2017; Falkner et al., 2023; Jung et al., 2023). While the MCMC corrections yield theoretical guarantees, shooting methods involve expensive molecular dynamics (MD) simulations and need to balance high rejection rates with large changes in trajectories. One-way shooting methods sample paths efficiently but yield highly correlated samples. Two-way shooting methods, which we compare to in Sec. 5, are more expensive but typically sample diverse paths faster. Recent machine learning approaches (e.g. Plainer et al. (2023); Lelièvre et al. (2023)) aim to reduce the need for MD.

Finally, various related methods rely on iterative simulation of SDE in (10) during training to learn the control drift term. Yan et al. (2022); Holdijk et al. (2023) are motivated from the perspective of stochastic optimal control, while Das et al. (2021); Rose et al. (2021) develop actor-critic methods using closely-related ideas from soft reinforcement learning. The variational method in Das et al. (2019) optimizes the rate function quantifying the probability of the rare events, while Singh and Limmer (2023) solves the Kolmogorov backward equation to learn the Doob's h -transform. However, all of these methods may be inefficient if the desired terminal state is sampled infrequently.

5 Experiments

We investigate the capabilities of our approach across a variety of different settings. We first illustrate features of our method on toy potentials before continuing to real-world molecular systems, including a commonly-used benchmark system, alanine dipeptide, and a small protein, Chignolin. The code behind our method is available at https://github.com/plainerman/variational-doob . Before diving into the experiments, we introduce the evaluation procedure and baseline methods.

Evaluation metrics. In our evaluation, we emphasize two key quantities: accuracy and efficiency. Efficiency is evaluated by the number of calls to the potential energy function, which requires extensive computation and dominates the runtime of larger molecules. For accuracy, we evaluate the log-likelihood of each sampled path and the maximum energy point (saddle point/transition state) along each sampled path. A good method samples many probable paths (i.e., high log-likelihood) and an accurate transition state (i.e., small maximum energy). See App. D for further details.

Baselines. We compare our approach against the MCMC-based two-way shooting method with uniform point selection with variable or fixed length trajectories. We found that two-way shooting

A

B

(a) MCMC

A

B

(b) Ours

(a) Single Gaussian

(b) Mixture of Gaussians

produced the most diverse path ensembles among possible baselines, although the acceptance probability can be relatively low for systems dominated by diffusive dynamics (Brotzakis and Bolhuis, 2016) and might be improved by better shooting point selection. This baseline gives theoretical guarantees about the ensemble and thus can be considered as a proxy for the ground truth. In that sense, our goal is not to beat two-way shooting but to approximate it with fewer potential evaluations.

5.1 Synthetic Datasets

Müller-Brown Potential. The Müller-Brown potential is a popular benchmark to study transition path sampling between metastable states. It consists of three local minima, and we aim to sample transition paths connecting state A and state B with a circular state definition. In Fig. 2, we visualize the potential and the sampled paths and can see that the same ensemble is sampled for both our method and two-way shooting. Our method exhibits a slightly reduced variance for unlikely transitions. In Table 1, we can observe that MCMC-based methods require many potential evaluations to achieve a good result, which comes from the low acceptance rate (especially when fixing the lengths of trajectories). Our method requires fewer energy evaluations (1 million vs. 1 billion) while finding paths with similar energy and likelihood. Note that the likelihood for variable approaches has been omitted, as it is governed by the number of steps in the trajectory and cannot be compared directly.

Gaussian Mixture. We further consider a potential in which the states are separated by a symmetric high-energy barrier that allows for two distinct reaction channels. In Fig. 3, we observe that a single Gaussian path cannot model a system with multiple modes of transition paths. Nevertheless, this issue can be resolved using a mixture of Gaussian paths, with slightly increased computational cost.

The Case for Neural Networks. According to our empirical study, the neural network parameterization of the Gaussian distribution statistics µ t m | 0 ,T , Σ t m | 0 ,T is an invaluable part of our framework. As an ablation, we consider parameterizing µ t m | 0 ,T , Σ t m | 0 ,T as piecewise linear splines whose intermediate points are updated using the same gradient-based optimizer as used for neural network training. In App. D.3, we report results comparing the W1 distance of learned marginals using neural network versus spline parameterizations, observing that splines yield inferior results even after an order of magnitude more potential function evaluations. We thus conclude that spline parameterizations are not competitive for learning transition paths and continue to focus on our neural-network approach.

5.2 Second-order Dynamics and Molecular Systems

Experiment Setup. We evaluate our methods on real-world high-dimensional molecular systems governed by the second-order dynamics (4): alanine dipeptide and Chignolin . Alanine dipeptide is a well-studied system of 22 atoms (66 total degrees of freedom), where the molecule can be described by two collective variables (CV): the dihedral angles ϕ , ψ . Chignolin is a larger system consisting of 10 residues with 138 atoms (414 total degrees of freedom) that cannot be summarized as easily. We use an AMBER14 force field (Maier et al., 2015) implemented in OpenMM (Eastman et al., 2017) but use DMFF (Wang et al., 2023) to backpropagate through the energy evaluations.

U(x)

T

S

SS

3000

2000

1000

0

200

400

600

800

1000

Alanine Dipeptide. In Table 2, we report results for four variants of our models, which either predict Cartesian coordinates or internal coordinates in the form of bond lengths and dihedral angles (compare App. D.4), either with or without Gaussian mixture. For our method, operating in internal coordinates yields better results compared to Cartesian coordinates, where the internal coordinates are distributed similarly to Gaussians and our network does need not learn equivariances (Du et al., 2022). Similarly, Gaussian mixture paths perform slightly better than a single Gaussian path due to the additional expressiveness. We note that paths sampled with Gaussian mixture exhibit a larger variance in max energy as they represent multiple reaction channels.

We find that prior-informed definitions of the desired initial and target states (i.e., CV) are necessary for MCMC to work efficiently with fixed-length trajectories. Finding these CVs in practice is challenging and only possible in this instance because the molecule is small and well-studied. For the larger system size in Table 2, it becomes intractable to use MCMC to connect precise states A,B ('exact') instead of larger regions ('relaxed'), even with a single trajectory. Variable length MCMC with relaxed endpoint conditions with CV perform well on this task, but our method is competitive using fewer evaluations and more strict boundary conditions. Fixed-length MCMC, even with prior-informed knowledge, can only find 100 trajectories while needing 50 times more potential evaluations compared to variable length.

Chignolin. The folding dynamics of Chignolin already pose a challenge and have not yet been well-studied compared to alanine dipeptide. We illustrate the qualitative experimental results for this system in Fig. 4. Operating in Cartesian space, our model samples a feasible transition within 25.6M potential energy evaluation calls and a transition with a duration of T = 1 ps .

6 Conclusion, Limitations and Future Work

In this paper, we propose an efficient computational framework for transition path sampling with Brownian dynamics. We formulate the transition path sampling problem by using Doob's h -transform to condition a reference stochastic process, and propose a variational formulation for efficient optimization. Specifically, we propose a simulation-free training objective and model parameterization that imposes boundary conditions as hard constraints. We compare our method with MCMC-based baselines and show comparable accuracy with lower computational costs on both synthetic datasets and real-world molecular systems. Our method is currently limited by rigidly defining states A and B to be a point mass with Gaussian noise instead of any arbitrary set. Finally, our method might be improved by accommodating variable length paths.

Acknowledgments and Disclosure of Funding

The authors would like to thank Juno Nam and Soojung Yang for spotting the unphysical steric barrier in the original pair of initial and target states, Jungyoon Lee for spotting an error in the energy computation, and Guan-Horng Liu, Maurice Weiler, Hannes Stärk and Yanze Wang for helpful discussions. The work of Yuanqi Du and Carla P. Gomes was supported by the Eric and Wendy Schmidt AI in Science Postdoctoral Fellowship, a Schmidt Futures program; the National Science Foundation (NSF), the Air Force Office of Scientific Research (AFOSR); the Department of Energy; and the Toyota Research Institute (TRI). The work of Michael Plainer was supported by the Konrad Zuse School of Excellence in Learning and Intelligent Systems (ELIZA) through the DAAD program Konrad Zuse Schools of Excellence in Artificial Intelligence, sponsored by the German Ministry of Education and Research and by Max Planck Society. The work of Kirill Neklyudov was supported by IVADO.

References

Anderson, J. B. (2007). Predicting rare events in molecular dynamics. Advances in Chemical Physics , 91:381-431.

Batatia, I., Benner, P., Chiang, Y., Elena, A. M., Kovács, D. P., Riebesell, J., Advincula, X. R., Asta, M., Baldwin, W. J., Bernstein, N., et al. (2023). A foundation model for atomistic materials chemistry. arXiv preprint arXiv:2401.00096 .

Blau, S. M., Patel, H. D., Spotte-Smith, E. W. C., Xie, X., Dwaraknath, S., and Persson, K. A. (2021). A chemically consistent graph architecture for massive reaction networks applied to solid-electrolyte interphase formation. Chemical science , 12(13):4931-4939.

Bolhuis, P. G. and Swenson, D. W. H. (2021). Transition path sampling as markov chain monte carlo of trajectories: Recent algorithms, software, applications, and future outlook. Advanced Theory and Simulations , 4(4).

Borrero, E. and Dellago, C. (2016). Avoiding traps in trajectory space: Metadynamics enhanced transition path sampling. The European Physical Journal Special Topics , 225(8-9):1609-1620.

Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.

Brotzakis, Z. F. and Bolhuis, P. G. (2016). A one-way shooting algorithm for transition path sampling of asymmetric barriers. The Journal of Chemical Physics , 145(16):164112.

Castellan, G. W. (1983). Physical Chemistry . Addison-Wesley, Reading, Mass, 3rd ed edition.

Chen, T., Liu, G.-H., and Theodorou, E. (2021a). Likelihood training of schrödinger bridge using forward-backward sdes theory. In International Conference on Learning Representations .

Chen, Y., Georgiou, T. T., and Pavon, M. (2016). On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications , 169:671-691.

Chen, Y., Georgiou, T. T., and Pavon, M. (2021b). Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. Siam Review , 63(2):249-313.

Chopin, N., Fulop, A., Heng, J., and Thiery, A. H. (2023). Computational doob h-transforms for online filtering of discretely observed diffusions. In International Conference on Machine Learning , pages 5904-5923. PMLR.

Das, A., Kuznets-Speck, B., and Limmer, D. T. (2022). Direct evaluation of rare events in active matter from variational path sampling. Physical Review Letters , 128(2):028005.

Das, A., Limmer, D. T., and Limmer, D. T. (2019). Variational control forces for enhanced sampling of nonequilibrium molecular dynamics simulations. The Journal of chemical physics , 151(24).

Das, A., Rose, D. C., Garrahan, J. P., and Limmer, D. T. (2021). Reinforcement learning of rare diffusive dynamics. The Journal of Chemical Physics , 155(13).

De Bortoli, V., Liu, G.-H., Chen, T., Theodorou, E. A., and Nie, W. (2023). Augmented bridge matching. arXiv preprint arXiv:2311.06978 .

De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. (2021). Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems , 34:17695-17709.

Dellago, C., Bolhuis, P. G., and Geissler, P. L. (2002). Transition path sampling. Advances in chemical physics , 123:1-78.

Dickson, A. (2018). Mapping the ligand binding landscape. Biophysical Journal , 115(9):1707-1719.

Doob, J. L. (1957). Conditional brownian motion and the boundary limits of harmonic functions. Bulletin de la Société Mathématique de France , 85:431-458.

Du, W., Zhang, H., Du, Y., Meng, Q., Chen, W., Zheng, N., Shao, B., and Liu, T.-Y. (2022). Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning , pages 5583-5608. PMLR.

Dürr, D. and Bach, A. (1978). The onsager-machlup function as lagrangian for the most probable path of a diffusion process. Communications in Mathematical Physics , 60:153-170.

Eastman, P., Swails, J., Chodera, J. D., McGibbon, R. T., Zhao, Y., Beauchamp, K. A., Wang, L.-P., Simmonett, A. C., Harrigan, M. P., Stern, C. D., Wiewiora, R. P., Brooks, B. R., and Pande, V. S. (2017). OpenMM 7: Rapid development of high performance algorithms for molecular dynamics. PLOS Computational Biology , 13(7):e1005659.

Escobedo, F. A., Borrero, E. E., and Araque, J. C. (2009). Transition path sampling and forward flux sampling. applications to biological systems. Journal of Physics: Condensed Matter , 21(33):333101.

Falkner, S., Coretti, A., Romano, S., Geissler, P., and Dellago, C. (2023). Conditioning normalizing flows for rare event sampling. arXiv preprint arXiv:2207.14530 .

Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V ., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. (2021). Pot: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8.

Freidlin, M. and Wentzell, A. (1998). Random perturbations of Dynamical Systems . Springer.

Gabrié, M., Rotskoff, G. M., and Vanden-Eijnden, E. (2022). Adaptive monte carlo augmented with normalizing flows. Proceedings of the National Academy of Sciences , 119(10):e2109420119.

Heng, J., De Bortoli, V., Doucet, A., and Thornton, J. (2021). Simulating diffusion bridges with score matching. arXiv preprint arXiv:2111.07243 .

Holdijk, L., Du, Y., Hooft, F., Jaini, P., Ensing, B., and Welling, M. (2023). Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems , 36.

Jamison, B. (1975). The markov processes of schrödinger. Zeitschrift für Wahrscheinlichkeitstheorie und verwandte Gebiete , 32(4):323-331.

Jang, E., Gu, S., and Poole, B. (2017). Categorical reparametrization with gumble-softmax. In International Conference on Learning Representations (ICLR 2017) . OpenReview. net.

Jung, H., Covino, R., Arjun, A., Leitold, C., Dellago, C., Bolhuis, P. G., and Hummer, G. (2023). Machine-guided path sampling to discover mechanisms of molecular self-organization. Nature Computational Science , 3(4):334-345.

Jung, H., ichi Okazaki, K., and Hummer, G. (2017). Transition path sampling of rare events by shooting from the top. The Journal of Chemical Physics , 147(15).

Juraszek, J. and Bolhuis, P. G. (2008). Rate constant and reaction coordinate of trp-cage folding in explicit water. Biophysical Journal , 95(9):4246-4257.

Kingma, D. P. and Welling, M. (2013). Auto-encoding variational bayes. International Conference on Learning Representations .

Kirmizialtin, S., Johnson, K. A., and Elber, R. (2015). Enzyme selectivity of HIV reverse transcriptase: Conformations, ligands, and free energy partition. The Journal of Physical Chemistry B , 119(35):11513-11526.

Kirmizialtin, S., Nguyen, V., Johnson, K. A., and Elber, R. (2012). How conformational dynamics of DNApolymerase select correct substrates: Experiments and simulations. Structure , 20(4):618-627.

Klucznik, T., Syntrivanis, L.-D., Ba´ s, S., Mikulak-Klucznik, B., Moskal, M., Szymku´ c, S., Mlynarski, J., Gadina, L., Beker, W., Burke, M. D., et al. (2024). Computational prediction of complex cationic rearrangement outcomes. Nature , 625(7995):508-515.

Koehl, P. and Orland, H. (2022). Sampling constrained stochastic trajectories using brownian bridges. The Journal of Chemical Physics , 157(5).

Lelièvre, T., Robin, G., Sekkat, I., Stoltz, G., and Cardoso, G. V. (2023). Generative methods for sampling transition paths in molecular dynamics. ESAIM: Proceedings and Surveys , 73:238-256.

Léonard, C. (2014). A survey of the schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems-A , 34(4):1533-1574.

Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2022). Flow matching for generative modeling. International Conference on Learning Representations .

Liu, G.-H., Lipman, Y., Nickel, M., Karrer, B., Theodorou, E., and Chen, R. T. (2023a). Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations .

Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E., Nie, W., and Anandkumar, A. (2023b). I 2 sb: Image-to-image schrödinger bridge. arXiv preprint arXiv:2302.05872 .

Liu, X., Wu, L., Ye, M., and Liu, Q. (2022). Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208.14699 .

Liu, X., Wu, L., Ye, M., and qiang liu (2023c). Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations .

Lu, Y., Stuart, A., and Weber, H. (2017). Gaussian approximations for transition paths in brownian dynamics. SIAM Journal on Mathematical Analysis , 49(4):3005-3047.

Maddison, C. J., Mnih, A., and Teh, Y. W. (2016). The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations .

Maier, J. A., Martinez, C., Kasavajhala, K., Wickstrom, L., Hauser, K. E., and Simmerling, C. (2015). ff14sb: improving the accuracy of protein side chain and backbone parameters from ff99sb. Journal of chemical theory and computation , 11(8):3696-3713.

Mehdi, S., Smith, Z., Herron, L., Zou, Z., and Tiwary, P. (2024). Enhanced sampling with machine learning. Annual Review of Physical Chemistry , 75.

Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023). Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning .

Neklyudov, K., Brekelmans, R., Tong, A., Atanackovic, L., Liu, Q., and Makhzani, A. (2024). A computational framework for solving wasserstein lagrangian flows. International Conference on Machine Learning .

Noé, F., Olsson, S., Köhler, J., and Wu, H. (2019). Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science , 365(6457):eaaw1147.

Noé, F., Schütte, C., Vanden-Eijnden, E., Reich, L., and Weikl, T. R. (2009). Constructing the equilibrium ensemble of folding pathways from short off-equilibrium simulations. Proceedings of the National Academy of Sciences , 106(45):19011-19016.

Onsager, L. and Machlup, S. (1953). Fluctuations and irreversible processes. Physical Review , 91(6):1505.

Papaspiliopoulos, O. and Roberts, G. (2012). Importance sampling techniques for estimation of diffusion models. Statistical methods for stochastic differential equations , 124:311-340.

Peluchetti, S. (2021). Non-denoising forward-time diffusions.

Peluchetti, S. (2023). Diffusion bridge mixture transports, Schrödinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00917 .

Plainer, M., Stärk, H., Bunne, C., and Günnemann, S. (2023). Transition path sampling with boltzmann generator-based MCMC moves. In Generative AI and Biology Workshop .

Rezende, D. J., Mohamed, S., and Wierstra, D. (2014). Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning , pages 1278-1286. PMLR.

Rose, D. C., Mair, J. F., and Garrahan, J. P. (2021). A reinforcement learning approach to rare trajectory sampling. New Journal of Physics , 23(1):013013.

Rotskoff, G. M. (2024). Sampling thermodynamic ensembles of molecular systems with generative neural networks: Will integrating physics-based models close the generalization gap? Current Opinion in Solid State and Materials Science , 30:101158.

Särkkä, S. and Solin, A. (2019). Applied stochastic differential equations , volume 10. Cambridge University Press.

Schauer, M., van der Meulen, F., and van Zanten, H. (2017). Guided proposals for simulating multi-dimensional diffusion bridges. Bernoulli , 23(4A).

Schrödinger, E. (1932). Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quantique. In Annales de l'institut Henri Poincaré , volume 2, pages 269-310.

Selli, D., Boulfelfel, S. E., Schapotschnikow, P., Donadio, D., and Leoni, S. (2016). Hierarchical thermoelectrics: crystal grain boundaries as scalable phonon scatterers. Nanoscale , 8(6):37293738.

Sharma, N. D., Singh, J., Vijay, A., Samanta, K., Dogra, S., and Bandyopadhyay, A. K. (2016). Pressure-induced structural transition trends in nanocrystalline rare-earth sesquioxides: A raman investigation. The Journal of Physical Chemistry C , 120(21):11679-11689.

Sheppard, D., Terrell, R., and Henkelman, G. (2008). Optimization methods for finding minimum energy paths. The Journal of chemical physics , 128(13).

Shi, Y., De Bortoli, V ., Campbell, A., and Doucet, A. (2023). Diffusion schrödinger bridge matching. arXiv preprint arXiv:2303.16852 .

Shoghi, N., Kolluru, A., Kitchin, J. R., Ulissi, Z. W., Zitnick, C. L., and Wood, B. M. (2023). From molecules to materials: Pre-training large generalizable models for atomic property prediction. In The Twelfth International Conference on Learning Representations .

Sidky, H., Chen, W., and Ferguson, A. L. (2020). Molecular latent space simulators. Chemical Science , 11(35):9459-9467.

Singh, A. N. and Limmer, D. T. (2023). Variational deep learning of equilibrium transition path ensembles. The Journal of Chemical Physics , 159(2).

Smith, J. S., Isayev, O., and Roitberg, A. E. (2017). Ani-1: an extensible neural network potential with dft accuracy at force field computational cost. Chemical science , 8(4):3192-3203.

Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. (2023). Aligned diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence , pages 1985-1995. PMLR.

Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. (2023). Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482 , 2(3).

Vanden-Eijnden, E. and Heymann, M. (2008). The geometric minimum action method for computing minimum energy paths. The Journal of chemical physics , 128(6).

Wang, H., Zhang, L., Han, J., and Weinan, E. (2018). Deepmd-kit: A deep learning package for manybody potential energy representation and molecular dynamics. Computer Physics Communications , 228:178-184.

Wang, X., Li, J., Yang, L., Chen, F., Wang, Y., Chang, J., Chen, J., Feng, W., Zhang, L., and Yu, K. (2023). DMFF: An open-source automatic differentiable platform for molecular force field development and molecular dynamics simulation. Journal of Chemical Theory and Computation , 19(17):5897-5909.

Weinan, E. and Vanden-Eijnden, E. (2010). Transition-path theory and path-finding algorithms for the study of rare events. Annual review of physical chemistry , 61(2010):391-420.

Wu, L., Gong, C., Liu, X., Ye, M., and Liu, Q. (2022). Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems , 35:36533-36545.

Xi, L., Shah, M., and Trout, B. L. (2013). Hopping of water in a glassy polymer studied via transition path sampling and likelihood maximization. The Journal of Physical Chemistry B , 117(13):3634-3647.

Yan, J., Touchette, H., and Rotskoff, G. M. (2022). Learning nonequilibrium control forces to characterize dynamical phase transitions. Physical Review E , 105(2):024115.

Zeng, J., Cao, L., Xu, M., Zhu, T., and Zhang, J. Z. (2020). Complex reaction processes in combustion unraveled by neural network-based molecular dynamics simulation. Nature communications , 11(1):5713.

Zhang, D., Bi, H., Dai, F.-Z., Jiang, W., Zhang, L., and Wang, H. (2022). Dpa-1: Pretraining of attention-based deep potential model for molecular simulation. arXiv preprint arXiv:2208.08236 .

A Proofs

A.1 Proofs from Sec. 2.2 (Doob's h -Transform Background)

Proposition 1. [Jamison (1975, Thm. 2)] Let h B ( x, t ) := ρ T ( x T ∈ B| x t = x ) denote the conditional transition probability of the reference process in (5) . Then,

where we use x t | T to denote a conditional process. The SDE in (6) is associated with the following transition probabilities for s < t < T ,

Note that all of our subsequent results hold for the case when B is a point-mass, with the only change being that the h -function becomes a density, h B ( x, t ) = ρ T ( B | x t = x ) .

Proof. See Jamison (1975) for a simple proof based on Ito's Lemma, assuming smoothness and strict positivity of h .

Proposition 2. The following PDEs are obeyed by (a) the density of the conditioned process ρ t | 0 ,T ( x ) := ρ t ( x | x 0 = A,x T ∈ B ) and (b) the h -function h B ( x, t ) ,

Reparameterizing (8b) in terms of s B ( x, t ) := log h B ( x, t ) , we can also write

Proof. Let p ( x t + s = y | x t = x ) denote the transition probability of a reference diffusion process

where ( G t ) ij = 1 2 Ξ t + s Ξ T t + s .

Now we condition the process on the end-point value x T ∈ B , and we get another kernel, i.e.

We let h B ( x, t ) = p ( x T ∈ B | x t = x ) denote the conditional probability over the desired endpoint condition given x t = x . According to laws of conditional probability, we can describe how h B ( x, t ) changes in time using the unconditioned transition probability

we take the derivative ∂ ∂s on both sides, and we get

Using the FP equation for the transition probability and integrating by parts, we have

Note that this holds ∀ x , hence, we have

without any loss of generality we can set t = 0

as desired to prove the optimality condition in (8b).

To prove (8a), denote p ( y, s ) = p ( x s = y | x 0 = x ) and differentiate p ( x s = y | x 0 = x, x T ∈ B ) = h B ( y,s ) h B ( x, 0) p ( y, s ) as

Note that h B ( x, 0) can be pulled outside the differential operator because it is a function of x . The PDE for the new kernel p ( y, s | B ) = p ( x s = y | x 0 = x, x T ∈ B ) (conditioned on the end-point) becomes

which matches the desired PDE in (8a) thereby proving the first two parts of Prop. 2.

Finally, to show (8c), we index time using t in Eq. (22) and change variables h B ( x, t ) = e s ( x,t ) ,

Next, we simplify 〈 ∇ , G t ∇ e s ( x,t ) 〉 = 〈 ∇ , G t e s ( x,t ) ∇ s ( x, t ) 〉 = 〈 ∇ e s ( x,t ) , G t ∇ s ( x, t ) 〉 + e s ( x,t ) 〈 ∇ , G t ∇ s ( x, t ) 〉 = e s ( x,t ) 〈 ∇ s ( x, t ) , G t ∇ s ( x, t ) 〉 + e s ( x,t ) 〈 ∇ , G t ∇ s ( x, t ) 〉 to finally write

which demonstrates (8c) since the inner term must be zero.

A.2 Proofs from Sec. 3.1 (Lagrangian Action Minimization for Doob's h -Transform)

We begin by proving Cor. 1, whose proof actually contains the initial steps needed to prove our main theorem Thm. 1. In both proofs, we omit conditioning notation q t ← q t | 0 ,T for simplicity and assume q t ( x ) s t ( x ) → 0 vanishes at the boundary x →±∞ , which is used when integrating by parts in x . Corollary 1. The Lagrangian objective in Thm. 1 which solves Doob's h -transform is equivalent to

Proof. Consider the following action functional

q

0

(

x

) =

δ

(

x

-

A

)

,

q

1

(

x

) =

The Lagrangian of this optimization problem is

where s t is the dual variable and we omit the optimization arguments, with S = min q t ,v t max s t L . Swapping the order of optimizations under strong duality, we take the variation with respect to v t in an arbitrary direction h t . Using G t = G T t , we obtain

Substituting into the above, we have

Integrating by parts in t and in x , assuming that q t ( x ) s t ( x ) → 0 as x →±∞ , yields

where in the second line, we integrate by parts in x again. Enforcing q T ( x ) = δ ( x -B ) and q 0 ( x ) = δ ( x -A ) and recalling S = min q t max s t L after eliminating v t , we recover the optimization in the statement of the corollary.

Theorem. 1. The following Lagrangian action functional has a unique solution which matches the Doob h -transform in Prop. 2,

δ

(

x

-

B

)

.

Namely, the optimal q ∗ t | 0 ,T ( x ) obeys (8a) and the optimal v ∗ t | 0 ,T ( x ) = ∇ x log h B ( x, t ) = ∇ x s ( x, t ) follows (8b) or (8c) .

Proof. The proof proceeds from (25) above,

We first show that the optimality condition with respect to s t yields the Fokker-Planck equation for q t in Prop. 2 (8a), before deriving the PDE in (8b) as the optimality condition with respect to q t .

Optimality Condition for (27) recovers Prop. 2 (8a) : The variation with respect to s t of (28) is simple, apart from the intermediate term. For a perturbation direction h t , we seek

where · indicates the functional on the right hand side. Proceeding to differentiate with respect to ε , we use linearity to pull d dε inside the integral and apply it first to obtain d dε ( s t + ε h t ) = h t . Using the product rule, recognizing the symmetry of terms, and evaluating at ε = 0 , we are left with

where in ( i ) we integrate by parts x .

We are now ready to set the variation of (28) with respect to s t (in an arbitrary direction h t ) equal to zero. Using (29), we have

which matches the desired optimality condition for the conditioned process in Prop. 2 (8a).

Optimality Condition for (27) recovers Prop. 2 (8b) : Starting again from (28), we take the variation with respect to q t . First, we repeat identical steps (integrate by parts in both x and t ) to reach (26),

where it is now clear that taking the variation with respect to q t and setting equal to zero yields

which is the desired PDE for s ( x, t ) = log h B ( x, t ) in (8c). To obtain (8b), we note an identity used to simplify the last term

Now, substituting s ( x, t ) = log h B ( x, t ) into Eq. (31) and abbreviating log h B ( · , t ) = log h t ( · ) , we obtain

which matches (8b) as desired.

The last equation defines the backward Kolmogorov equation for the diffusion process with the drift b t ( x ) and covariance matrix G t , i.e. the function h t ( x ) defines the conditional density h t ( x ) = p ( x T ∈ B ′ | x t = x ) for some set B ′ , which agrees with the forward process with the same drift and covariance. The boundary condition q T ( x ) = δ ( x -B ) together with the backward equation define the unique solution to this PDE. Since the PDEs and the boundary conditions are the same as in Doob's h -transform, we have h t ( x ) = p ( x T = B | x t = x ) .

Corollary 2. The following SB problem

yields the path measure P ∗ 0: T associated with the SDE in (6) as its unique minimizing argument. The temporal marginals of P ∗ 0: T are equal to those which optimize the Lagrangian objective in Thm. 1.

Proof. Weuse the Girsanov theorem (Särkkä and Solin, 2019, Sec. 7.3) to calculate the KL divergence between the following two Brownian diffusions with fixed initial condition x 0 = A ,

In particular, noting the difference of drifts is b t ( x t ) + 2 G t v t | 0 ,T ( x t ) -b t ( x t ) = 2 G t v t | 0 ,T ( x t ) , the likelihood ratio is given by

We finally calculate the KL divergence, noting that, after taking the log, the expectation of the integral ∫ ( · ) dW t in the final term vanishes,

which matches (9a) up to a constant factor of 2 does not change the optimum. We finally compare to the constraints in Thm. 1. First, it is clear that the diffusion in (34) satisfies the Fokker-Planck equation in (9b) (Särkkä and Solin, 2019, Sec. 5.2). We respect (9c) by optimizing over endpoint-constrained path measures, which yields

as desired.

B Gaussian Path Parameterizations

Proposition 3. For the family of endpoint-conditioned marginals q t | 0 ,T ( x ) = N ( x | µ t | 0 ,T , Σ t | 0 ,T ) ,

satisfies the Fokker-Planck equation (12) for q t | 0 ,T and diffusion coefficients G t = 1 2 Ξ t Ξ T t .

Proof. Consider the following identities for the Gaussian family of marginals q t ( x ) = N ( x | µ t , Σ t ) , where we omit conditioning q t ← q t | 0 ,T for simplicity of notation,

We begin by solving for a vector field u o t ( x ) that satisfies the continuity equation (where u o t denotes the drift of an ODE)

The vector field satisfying this equation is

which we can confirm using the identities in (38). Indeed, for the terms on the RHS of Eq. (39),

Putting these terms and the time derivative from (38c) into Eq. (39) we conclude the proof.

However, we are eventually interested in finding the formula for the drift u t that satisfies the FokkerPlanck equation in (12). That is, to describe the same evolution of density ∂q t ( x ) ∂t , the relationship between u t and u o t is as follows

t

Finally, we use the identities in (38) to obtain

Proposition 4. Given a set of processes q k t | 0 ,T ( x ) and mixtures weights w k , the vector field satisfying the Fokker-Planck equation in (12) for the mixture q t | 0 ,T ( x ) = ∑ k w k q k t | 0 ,T ( x ) is given by

where u ( q,k ) t | 0 ,T ( x ) satisfies the Fokker-Planck equation in (12) for q k t | 0 ,T ( x ) . This identity holds for both first-order dynamics in spatial coordinates only or second-order dynamics in x = (¯ x, ¯ v ) .

Proof. See Peluchetti (2023) Theorem 1 and its proof in their App. A.

C Extended Related Work

C.1 Machine Learning for Molecular Simulation

The main dilemma of molecular dynamics comes from the accuracy and efficiency trade-off-accurate simulation requires solving the Schrödinger equation which is computationally intractable for large systems, while efficient simulation relies on empirical force fields which is inaccurate. Recently, there has been a surge of work in applying machine learning approaches to accelerate molecular simulation. One successful paradigm is machine learning force field (MLFF) which leverages the transferability and efficiency of machine learning methods to fit force/energy prediction models on quantum mechanical data and transfer across different atomic systems Smith et al. (2017); Wang et al. (2018). More recently, increasing attention has been focused on building atomic foundation models to encompass all types of molecular structures Batatia et al. (2023); Shoghi et al. (2023); Zhang et al. (2022).

Sampling is a classical problem in molecular dynamics to draw samples from the Boltzmann distribution of molecular systems. Classical methods mainly rely on Markov chain Monte Carlo (MCMC) or MDwhich requires long mixing time for multimodal distributions with high energy barriers Rotskoff (2024). Generative models in machine learning demonstrate promises in alleviating this problem by learning to draw independent samples from the Boltzmann distribution of molecular systems (known as Boltzmann generator) Noé et al. (2019). Numerous methods have been developed to utilize generative models as a proposal distribution for escaping local minima in running MCMC methods Gabrié et al. (2022). However, one critical issue is that generative models rely on training from samples. Although recent advances have been developed to learn from unnormalized density (i.e., energy) function, the training inefficiency limits their applicability to solve high-dimensional molecular dynamics problems. To circumvent the curse of dimensionality for the sampling problem, another branch of work study to learn coarse-grained representation with neural networks Sidky et al. (2020). For broader literature of applying machine learning to enhanced sampling, we refer the reader to Mehdi et al. (2024).

D Further Experimental Details

D.1 Evaluation Metrics

To assess the quality of our approach in terms of performance and physicalness of paths, we compare them under different metrics to well-established TPS techniques. One important describing factor of a trajectory is the molecule's highest energy during the transition. These high-energy states are often referred as transition states and less likely to occur but they determine importance factors during chemical reaction such as reaction rate. As such, we will look at the maximum energy along the transition path and use it to compare the ensemble of trajectories more efficiently. The main goal is to show that lower energy of the transition states can be sampled by the methods.

However, the maximum energy does not account for the fact that the transition path needs to be sequential, and each step needs to be coherent based on the previous position and momentum. For this, we also compare the likelihood of the paths (i.e., unnormalized density) by computing the probably of being in the start state ρ ( x 0 ) and multiplying it with the step probability such that

For the step probability π , we solve the Langevin leap-frog implementation as implemented in OpenMM to solve N ( x i +1 | x i + dt · b t i ( x ) , dtσ 2 i ) . As for the starting probability, we compute the unnormalized density of the Boltzmann distribution for our start state z and assume that the velocity v can be sampled independently (Castellan, 1983, Sec. 4.6)

with the Boltzmann constant k B and the diagonal matrix M containing the mass of each atom.

As for the performance, the number of energy evaluations will be the main determining factor of the runtime for larger molecular systems, especially for proteins. We hence compare the use of the number of energy computations as a proxy for hardware-independent relative measurements. In our tests, this number aligned with the relative runtime of these approaches.

D.2 Toy Potentials

The toy systems move according to the following integration scheme (first-order Euler)

following the definition of our stochastic system in Sec. 2.2 with a time-independent Wiener process, where ξ is a constant time-independent standard deviation for all dimensions.

Müller-Brown. The underlying Müller-Brown potential that has been used for our experiments can be written as

We used a first-order Euler integration scheme to simulate transition paths with 275 steps and a dt of 10 -4 s . ξ was chosen to be 5 and 1,000 transition paths were simulated. We have used an MLP with four layers and a hidden dimension of 128 each, with swish activations. It has been trained for 2,500 steps with a batch size of 512 and a single Gaussian.

In Fig. 5a, we compare the likelihood of the sampled paths. We can see that one-way shooting takes time until the path is decorrelated from the initial trajectory, which is shorter and thus has a higher likelihood. All MCMC methods exhibit this behavior, which is typically alleviated by using a warmup period in which all paths are discarded. After that, all methods exhibit similar likelihood, with our method having a slightly lower likelihood. Looking at the transition state (i.e., maximum energy on the trajectory) in Fig. 5b reveals that all methods have a similar quality of paths.

0

200

400

600

800

1000

10 3

8×10 2

9×10 2

log path likelihood

one-way-shooting

two-way-shooting

var-doobs

0

200

400

600

800

1000

1.2 × 10 2

1.4 × 10 2

1.6 × 10 2

1.8 × 10 2

2×10 2

2.2 × 10 2

2.4 × 10 2

Maximum energy

one-way-shooting

two-way-shooting

var-doobs

(b) Maximum Energy

We can further analyze the quality of our method by investigating the difference between the 'ground truth' marginal ρ t | 0 ,T ( x ) and the learned marginal q t | 0 ,T ( x ) . For this, we compute the Wasserstein W1 distance (Flamary et al., 2021) between the marginal observed by fixed-length two-way shooting (which we assume to be close to the ground truth) and our variational approach. We observe a mean W1 distance of 0 . 130 ± 0 . 026 and visualize it along the time coordinate t (in steps) in Fig. 6.

0

50

100

150

200

250

0.04

0.06

0.08

0.10

0.12

0.14

0.16

Wasserstein W1 Distance

t

Dual-Channel Double-Well. To demonstrate the advantage of mixtures, we have used the twodimensional potential

In this case, we have used dt = 5 ∗ 10 -4 s with a transition time of T = 1 s and ξ = 0 . 1 . As for the MLP, we have used the same structure as in the Müller-Brown example but trained it for 20,000 iterations. The corresponding weights to Prop. 4 are w = [ 1 2 , 1 2 ] and are fixed for this experiment and hence w ̸∈ θ .

D.3 Neural Network Ablation Study

In Fig. 7 we compare how different parameterizations of µ ( θ ) t | 0 ,T , and Σ ( θ ) t | 0 ,T impact the quality of trajectories on the Müller-Brown potential. For this, we compare linear and cubic splines (with 20 knots) with neural networks. As a metric to estimate the quality, we compare the Wasserstein W1 distance between the learned marginal q t | 0 ,T ( x ) and the marginal observed by fixed-length two-way shooting (i.e., baseline). We notice that using linear splines results in the highest W1 distance, while cubic splines improve the quality. Using neural networks, however, yields the best approximation.

We have fixed the computational budget for all systems, which means that we have trained splines for more epochs than the neural network (since they are slower to train). For high-dimensional systems, the runtime is mostly determined by the number of potential evaluations and not the complexity of the architecture. We thus conclude that the additional expressivity provided by neural networks is necessary for more complicated (molecular) systems and does not introduce much computational overhead.

0

50

100

150

200

250

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

Wasserstein W1 Distance

Linear Spline

Neural Network

t

D.4 Molecular Systems

To simulate molecular dynamics, we rely on the AMBER14 force field (amber14/protein.ff14SB Maier et al. (2015)) without a solvent, as implemented in OpenMM (Eastman et al., 2017). As OpenMM does not support auto-differentiation, we do not use OpenMM for the simulations themselves, but utilize DMFF (Wang et al., 2023) which is a differentiable framework implemented in JAX (Bradbury et al., 2018) for molecular simulation. This is needed because during training we compute ∇ θ U ( x t | 0 ,T ∼ N ( µ ( θ ) t | 0 ,T , Σ ( θ ) t | 0 ,T ) ) , where the concrete x t | 0 ,T is sampled based on the parameters of the neural network.

For the concrete simulations, we ran them with the timestep dt = 1 fs , with T = 1 ps , γ = 1 ps , and Temp = 300 K . To compute the MCMC two-way shooting baselines, we use the same settings and consider trajectories as failed, if they exceed 2,000 steps without reaching the target.

Neural Network Parameterization. We parameterize our model with neural networks, a 5-layer MLP with ReLU activation function and 256/512 hidden units for alanine dipeptide and Chignolin, respectively. The neural networks are trained using an Adam optimizer with learning rate 10 -4 .

We represent the molecular system in two ways: (1) in Cartesian coordinates, which are the 3D coordinates of each atoms, and with (2) internal coordinate which instead uses bond length, angle and dihedral angle along the molecule, where we use the same parameterization as in (Noé et al., 2019).

Our state definition includes a variance parameter for the initial and target marginal distributions at t = 0 and t = T , we choose the variance to be 10 -8 which almost does not change the energy of the perturbed system.

Visualization of Transition for Alanine Dipeptide. In Fig. 8, we show a transition sampled without any noise from the model with internal coordinates and 2 Gaussian mixtures.

Comparison of Sampled Paths During Training. In our training procedure, the marginal starts with a linear interpolation between A and B , which produces very unlikely paths with potentially high energy states. In Fig. 9, we compare how the quality of paths changes depending on the number of training iterations (i.e., the number of potential evaluations). We show the curve for a single Gaussian mixture with Cartesian coordinates. Similar trends can be observed in other settings.

1M

10M

20M

30M

40M

Number of Potential Evaluations

10 3

10 4

Maximum Energy [kJ/mol]

Loss Curves. In this section, we would like to investigate the training losses of different configurations. For this, we plot the exponential moving average of the loss ( α = 0 . 001 ) to better highlight the trends of the noisy variational loss. Fig. 10 compares the results of different training settings. We can observe that mixtures can decrease the overall loss, but all model variations converge to a similar loss value.

0

20000

40000

60000

80000

100000

120000

140000

Epoch

10 6

10 7

10 8

10 9

10 10

10 11

10 12

10 13

Loss

Internal

Internal, 2 Mixtures

Cartesian

Cartesian, 2 Mixtures

D.5 Computational Resources

All our experiments involving training were conducted on a single NVIDIA A100 80GB. The baselines themselves were computed on a M3 Pro 12-core CPU.

E Societal Impact

Our research concerns the efficient sampling of transition paths which are crucial for a variety of tasks in biology, chemistry, materials science and engineering. Our research could potentially benefit research areas from combustion, catalysis, protein design to battery design. Nevertheless, we do not foresee special potential negative impacts to be discussed here.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The major claims made in the abstract and introduction are backed up by the theoretical and empirical results in the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Sec. 6, we discuss potential limitations and future work.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: A detailed set of assumptions is presented in Sec. 3 and further proofs for theoretical results are included in App. A.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We disclose all the information needed to reproduce the main experimental results and make our code available at https://github.com/plainerman/variational-doob .

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The data we use are synthetic datasets that can be generated by publicly available codes and real-world molecular systems taken from previous literature cited in the paper. Our code is publicly available at https://github.com/plainerman/variational-doob .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We report the experimental details in Sec. 5 and further specify the parameter and settings in App. D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We appropriately report the mean and standard deviation in the experiments where suitable (compare Table 1 and Table 2).

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We include computational resources used to run the model in App. D.5.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

Answer: [Yes]

Justification: Yes, we believe our research conforms with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the potential societal impact of our method in App. E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not believe that any safe guards needs to be in place, as our approach does not pose any direct risks. As for the data, only publicly available and commonly used resources have been used and, hence, does not pose a risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We use public datasets and libraries in this research and have cited the sources.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: [NA]

Guidelines: No new assets are introduced in this research.

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No crowdsourcing or human subjects are involved in this research.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: No human subjects are involved in this research.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The paper is concerned with sampling trajectories with a terminal condition. For stochastic processes governed by a Brownian Motion, Doob's h-transform gives an posterior SDE that leads to samples with the final condition. However, estimating the h-function that is needed for the posterior SDE usually involves simulating trajectories, which is inefficient if the terminal condition is rarely reached.
The authors propose a simulation-free variational optimization method to estimate the h-function based on a least action principle and Gaussian approximations to the marginal densities.

STRENGTHS:
1. The authors provide a solid and clear background on Doob's h transform that gets supported by the provided proofs in the appendix.

2. The paper clearly highlights the challenges of the optimization problem and proposes an efficient solution adressing these challenges.

3. The related work section gives a good overview and nicely connects to related topics.

WEAKNESSES:
1. I found the path histograms in Figure 2 to be too cluttered. I would propose adding less samples.

2. The paper misses a learning curve. It would in general be interesting to have more training details.

Minor Weaknesses:

In Chapter 3 and in the appendix, the authors change from trajectory length T to the unit interval. There should be a sentence that explains this change.

QUESTIONS:
I could not totally follow the derivation in the appendix. Can you explain how to get to Equation 25 from the previous equation in line 548?
I do not see why second term in equation 25 gets subtracted, while in line 548 only the last term has a minus.

LIMITATIONS:
The authors adequately adressed the limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-19

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a variational formulation of Doob's h-transform, which characterizes the distribution over paths with a given endpoint. Instead of relying on potentially wasteful sampling approaches, the authors propose directly optimizing a tractable variational distribution over transition paths which satisfy the initial and terminal conditions by design. This approach reduces the search space over trajectories and avoids trajectory simulation. Experiments on real-world molecular simulation and protein folding tasks demonstrate the applicability of the approach.\n\nSTRENGTHS:\nThe paper clearly motivates the problem it tackles, describes the challenges well and nicely introduces the idea behind their method in an illustrative fashion. Overall the paper is very well-written and structured. In particular, it builds up the method piece-by-piece explaining the choices along the way. As far as I can judge the related work seems to be exhaustive. The experiments are done on interesting problems as far as I can tell and I also appreciated that the authors made their code publicly available.\n\nWEAKNESSES:\nMy main concern comes from trying to interpret the experimental results and in particular judging the performance compared to MCMC. It seems quite clear from Tables 1 and 2 that the variational approach requires fewer calls to the potential energy function than MCMC, however their performance differences are harder to judge in my opinion. \nMore specifically, in Table 1 the standard deviations are so large that there is basically no meaningful statistical difference between the shown results especially for the Max Energy, but also the Log-likelihood. Now this might mean that MCMC and the presented method both perform well, but its surprising to me that there is that much inherent variation. Similarly in Table 2, the variance for the Max Energy of MCMC in the first line is huge. \nI'm also a bit confused by the Max Energy increases when using a mixture in Table 2.\nTo better understand how the performance of the variational approach improves during training, it would be nice to see a plot that shows the Max Energy as a function of the training epochs.\n\nFinally, I would have liked to see a discussion of the limitations of the approach. Section 6 has Limitations in the title, but does not actually discuss them in any way beyond extensions of the proposed method.\n\nQUESTIONS:\n- How do you choose the number of mixture components in practice?\n- Why does the MaxEnergy increase when using a mixture distribution as a variational approximation in Table 2?\n- How do you explain the huge variance of the Max Energy of MCMC (variable length) in the first row of Table 2?\n- What are the main limitations of the current approach?\n\nLIMITATIONS:\nSection 6 has \"Limitations\" in the title, but limitations are not discussed, only future work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe Authors of this paper tackle the problem of sampling conditioned SDEs with a specific interest in \"transition path sampling\", i.e. sampling a Langevin-type SDE undergoing a transition between an initial state (or set of states) $A$ and a final target set $B$. Sampling transition paths efficiently can provide a tremendous boost to research in catalysis or drug design. In this paper, the Authors model the transition path distribution with either (i) a parametrized Gaussian process or (ii) a parametrized mixture of Gaussians processes. These families of priors are then optimized by leveraging a variational formulation of the problem developed by the Authors starting from Doob's $h$-transform.\n\nSTRENGTHS:\nThe Authors tackle a challenging problem with a novel variational formulation, which is apt to be optimized by leveraging techniques developed by the generative modeling community within ML. A clear strength of the paper is the sound theoretical analysis justifying the proposed method. Interestingly, the choice of Gaussian process (or a mixture of Gaussian processes) variational priors allows the Authors to simplify the algorithm using analytical results and sidestep lengthy calculations at the deployment phase.\n\nWEAKNESSES:\nI find the paper overall clearly written, but I had to go through section 3.2 (Computational Approach) several times to grasp how the method can be deployed in practice. Specifically, I think that it can be improved the explanation of the fact that, because of the Gaussian prior, you only need to model the transition path probability and not the $h$-transform itself, as well as the explanation of the actual optimization step (Reparametrization of Gradients).\n\nThe subscript $0,T$ is introduced for the first time in Eq. (6) without explanation and used throughout the manuscript to label variables related to the conditioned process. It might be useful to clarify this notation explicitly. \n\nThe experimental evaluation is somewhat limited, especially concerning the Chignolin protein. Over the years, many transition path sampling strategies have been developed, as well as very much related enhanced sampling techniques. It would be nice to have a comparison also to different baselines, as well as a discussion on the computational complexity and actual running times of the baselines.\n\nQUESTIONS:\nIs the $h$-transform related to the committor function in some way? If so, how does your Thm 1 relates to the variational formulation of the committor (see e.g. Eq. 20 of \"Transition Path Theory and Path-Finding Algorithms for the Study of Rare Events\" by Weinan E and Eric Vanden-Eijnden)\n\nLIMITATIONS:\nLimited experimental evaluation", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submitted manuscript presents a variational formulation of Doob’s h-transform, leading to a novel (simulation-free) computational approach for rare event sampling in transition paths. The task of interest involves conditioning a dynamical system driven by Brownian motion with a known drift term to reach a given endpoint. In theory, this terminal condition can be addressed using Doob’s h-transform, resulting in an associated SDE for the conditional dynamic. However, practical implementation requires knowledge of the h-transform. The authors propose a variational problem formulation that provides the necessary information about the h-transform as its solution. Solving this variational problem results in a computational approach for simulating the conditional dynamic. Compared to existing methods, it is claimed that the proposed approach avoids importance sampling estimators and expensive trajectory simulations. The method is tested on both synthetic and real datasets.\n\nSTRENGTHS:\nThe paper introduces a novel variational objective to describe Doob’s h-transform, leading to a promising and innovative computational approach for simulating the conditional dynamics of transition paths. Although I appreciate the general approach presented in Section 3.1, I am not convinced by the computational approach proposed in Section 3.2 due to several reasons pointed out below. If these issues can be addressed, I believe that the approach can be further utilized to construct highly efficient sampling strategies.\n\nWEAKNESSES:\n- The proposed approach in Section 3.2.1 introduces certain issues. The connection to Doob’s h-transform established in Theorem 1 is only guaranteed if an exact solution to the optimization problem (9) or (11) is found. However, by introducing the Gaussian parametrization of $q_{t|0,T}$, this guarantee is lost. There is a lack of discussion regarding the implications and potential effects of this parametrization on the overall accuracy and validity of the method.\n- Moreover, the formulation of the task of solving equation (12) given $q_{t|0,T}$ is misleading. In reality, you are solving an inverse problem here. Given the solution of the PDE in (12), your goal is to reconstruct $u_{t|0,T}$, which is generally an ill-posed problem. For instance, introducing uncertainties into the model description via the drift $b_t$ or diffusion $\\Xi_t$ can lead to significant challenges. This scenario is likely when applying the approach in practice under model misspecification. A critical question to address is the robustness of the proposed framework against model misspecification, such as small perturbations in $b_t$.\n- The claim that the approach is simulation-free is not entirely clear. In line 173, it is stated that the Gaussian parametrization allows for the generation of arbitrary samples. However, this assumes that the states $x_{t|0,T}$ are independent for all $t$, which is not true when $x_{t|0,T}$ is a solution of the SDE in (10). The dependence between the states needs to be accounted for in the sampling process, and this aspect seems to be overlooked. \n- I am wondering what the measured quantities in the numerical experiments actually reveal about the correctness of the proposed approach. Evaluating the maximum energy and the likelihood might not provide sufficient information about the accuracy of the estimated distribution. If the sampled paths only follow high likelihood regions, this does not necessarily indicate that the correct distribution has been captured. The objective should be to estimate quantities of interest related to the conditioned SDE in (6). However, it is not clear whether the simulated paths correspond to accurate simulations of (6). For instance, how does the method perform when estimating rare event probabilities or other Monte Carlo estimators with respect to (6)? This aspect needs to be thoroughly addressed to validate the effectiveness of the proposed approach.\n\nQUESTIONS:\n- To enhance the practical applicability of the proposed framework, it is crucial to study the impact of inexactly solving the variational problems (9) or (11). Specifically, how do errors propagate to the conditioned dynamical system driven by the SDE (6) when an approximate solution of (9) or (11) is used? Understanding this error propagation is essential for assessing the robustness and reliability of the proposed method in practical scenarios.\n- Is the Optimization problem in (9) well-posed? This means, is the optimal solution unique?\n- In several instances, there are missing commas in the notation for the inner product (e.g., equations (9b) and (12)).\n- Many important mathematical details and assumptions are missing. Most importantly, what are the assumptions on the drift vector field $b_t$ to ensure the well-posedness of the proposed scheme?\n- Regarding Section 3.2.2: Instead of introducing $\\xi_{\\min}$ I suppose that one could also directly work with a pseudoinverse when defining $G_t^{-1}$.\n- In Proposition 4, should it be $u_{t|0,T}^{(k)}$ on the right hand side of the equation for $u$?\n- When referring to probability density functions, it's important to use proper notation. Instead of $\\rho(x_t=x)$, it would be clearer to use $\\rho_t(x)$ depending on the context.\n- In the abstract, you claim that no \"inefficient\" importance sampling estimators are required. However, it would be beneficial to compare this approach with such estimators to demonstrate its advantages more clearly.\n\nLIMITATIONS:\n- There is no discussion about error propagation arising due to the Gaussian parametrization.\n- The approach is limited to conditioning the sample path on point sets $x_T = B$.\n- Limited statistical evidence is presented in the numerical experiments.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a variational formulation of the Doob's h-transform to change the problem from the expensive simulations of trajectories to an optimization problem over possible trajectories from given initial point and end point. The model parameterization introduced imposes the desired boundary conditions by design and uses a mixture of Gaussian as the variational family. The variational formulation offers an alternative to expensive simulations and MCMC which is hard to scale. The proposed framework is applied on a simulated case study and real world problem in material science/computational chemistry of transition path sampling  which is a study of molecule transitions betweeb local energy minima and metastable states under random fluctuations. The results are along expected lines where it is possible to scale the solution on computationally intractable problems with MCMC and other sampling based methods.\n\nSTRENGTHS:\n1. The paper is well written, with sound math, a lot of references on recent work and older literature on the subject. I found the introduction to be a great summary of the work. \n2. The idea is sound and through the use of variational framework to solve a complex problem of estimating rare event probabilties by sampling forward trajectories brings out the desired qualities of sample efficiency, reducing search space and matching a complex target distribution with an easier to sample variational distribution. As also observed in Table1 and 2, MCMC cannot be scaled to certain large scale experiments where Variational inference can give good results. \n3. The work finds application for solving transition path sampling problem which is important in material science and chemistry from what it seems(I am not so well acquanited with those domains.) The case study on protein is well documented and explained. \n4. The figures and illustrations are clean and support the narrative.\n\nWEAKNESSES:\nMinor things\n- Increase font sixe for Table1, bold the important results.\n- I would have liked to see a bit more dicussion on the dimensionality, what are the typical values of D ..\n\nQUESTIONS:\n1. I am curious,  similar to Figure 3, the authors can show an experiment with familiy of dtistributions other than Gaussian. How does the expressivity then come to play ? \n2. as for any SDE problem solution, what is the effect of discretization values and parameters in practice as given in line 184-185.\n3. How big is the constraint of fixed length(T) on transition paths in practice as solved in this method ? \n4. What is the behaviour of the problem when there are multiple basins of attraction ?\n\nLIMITATIONS:\nThe authors have addressed the challenges quite well and what the future research directions could be taken.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}