Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Authors: Sanjeev Raja
Martin Sipka
Michael Psenka
Tobias Kreiman
Michal Pavelka
Aditi S. Krishnapriyan
Venue: ICML 2025 poster (2025)
Keywords: Transition path sampling
molecular dynamics
generative models
diffusion models
flow matching models
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We repurpose generative models for transition path sampling via minimization of the Onsager-Machlup action functional.

Abstract: Transition path sampling (TPS), which involves finding probable paths connecting two points on an energy landscape, remains a challenge due to the complexity of real-world atomistic systems. Current machine learning approaches rely on expensive training procedures and under-utilize growing quantities of atomistic data, limiting scalability and generalization. Generative models of atomistic conformational ensembles sample temporally independent states from energy landscapes, but their application to TPS remains mostly unexplored. In this work, we address TPS by interpreting candidate paths as trajectories sampled from stochastic dynamics induced by the learned score function of generative models, namely denoising diffusion and flow matching. Under these dynamics, finding high-likelihood transition paths becomes equivalent to minimizing the Onsager-Machlup (OM) action functional, enabling us to repurpose pre-trained generative models for TPS in a zero-shot fashion. We demonstrate our approach on a Müller-Brown potential and several fast-folding proteins, where we obtain diverse, physically realistic transition pathways, as well as tetrapeptides, where we demonstrate successful TPS on systems not seen by the generative model during training. Our method can be easily incorporated into new generative models, making it practically relevant as models continue to scale and improve.

Action-Minimizatiin Meets Generative Modeling:
Efficient Transition Path Sampling with the Onsager-Machlup Functional
Sanjeev Raja" "I Martin Sipka 22
Michael Psenka
Tobias Kreiman
Michal Pavelka Aditi S. Krishnapriyan 145
Abstract
statistical mechanics and computational science (Tuckerman, 2023; Frenkel & Smit, 2023). A key area of interest is the sampling of rare, transition events between two stable con- figurations, such as in chemical reactions or protein folding (Bolhuis et al., 2002; Dellago et al., 2002). This task, broadly known as transition path sampling (TPS), is challenging due to the presence of energy barriers, which create substantial difference in timescales between rare events and the fastest dynamical motions of the system (e.g., bond vibrations). This has inspired a rich line of literature on enhanced sampling techniques (Torrie & Valleau, 1977; Swendsen & Wang, 1986; Laio & Parrinello, 2002b: Laio & Gervasio, 2008; Tiwary & Parrinello, 2013; Valsson & Parrinello, 2014). More recently, machine learning (ML) based methods have gained popularity for accelerating TPS by learning a control drift term to bias system towards desired target state (Sipka et al., 2023; Holdijk et al., 2024; eu 2024; Seong et al., 2024). However, these approaches rely on highly spe- cialized training procedures and fail to exploit the growing quantity of atomistic simulation and structural data (Bank, 1971; Vander Meersche et al., 2024; Lewis et al., 2024), and the increasing availability of high-quality atomistic conformational generative models (Abramson al., 2024; Lewis et al., 2024; Jing et al., 2024a; Zheng et al., 2024).
Transition path sampling (TPS), which involves finding probable paths connecting two points on an energy landscape, remains a challenge due to the complexity of real-world atomistic systems. Current machine learning approaches use expensive, task-specific, and data-free training procedures, limiting their ability to benefit from re- cent advances in atomistic machine learning, such as high-quality datasets and large-scale pre-trained models. In this work, we address TPS by inter- preting candidate paths as trajectories sampled from stochastic dynamics induced by the learned score function of pre-trained generative models, specifically denoising diffusion and flow matching. Under these dynamics, finding high-likelihook transition paths becomes equivalent to minimizing (OM) action functional. This enables us to repurpose pre-trained generative models for TPS in a zero-shot manner, in contrast with bespoke, task- ses--specfiic TPS models trained in previous work. We demonstrate our approach varied molecular systems, obtaining diverse, physically realistic transition pathways and generalizing beyond the pre-trained model's 's original training dataset. Our method can be easily incorporated into new generative models, making it practically relevant as models continue to scale and improve with increased data availability.
the Onsaagr-Macchun
generattieee
Generative models can produce unbiased, independent samples from atomistic conformational ensembles (Noé et al., 2019; Zheng et al., 2024; Jing et al., 2024a) and have shown the potential to generalize across chemical space (Klein & Noé, 2024; Lewis et al., 2024). However, they have not been directly used for TPS due to the use of uncorrelated states during training. In this work, we propose conceptually simple post-training method to repurpose generative models to perform TPS in zero-shot manner. Our core idea exploits the fact that generative models based on denoising diffusion (Ho et al., 2020) and flow matching (Lipman al., 2022) induce set of stochastic Langevin dynamics on the data manifold, governed by their learned score function. Drawing inspiration from statistical mechanics, the probability of paths sampled from these dynamics can be characterized by quantity known as the Onsager-Machlup (OM) action functional (Onsager & Machlup, 1953). As result, we can identify high-probability paths between arbitrary points on the data manifold directly
1. Introduction
Efficiently sampling the configurational distribution of high-dimensional molecular systems is grand challenge in
"Equal contribution, random order. 'Dept. of EECS, UC Berkeley Faculty of Mathematics and Physics, Charles University "Done prior to joining Amazon 'Dept. of Chemical and Biomolec- ular Engineering, UC Berkeley 2 'Applied Mathematics and Computational Rescarch, LBNL. Correspondence to: Sanjeev Raja, Aditi S. Krishnapriyan < (sanjecvr, aditik } @berkeley.edur
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada, PMLR 267, 2025. Copyright 2025 by the author(s).

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Generative Models of Atomistic Distributions
OM Action (S) Optimization
Generative Process
Pitai
Piuaa
Po
B
Learned Vector Field Flow Matching: ugg-
Learned
X
()) Path Term
(II) Drift Term
(III) Divergence Term
Score Function
x
A:
DAf
Bo (xx
V B (x (0) rj
Diffusion: EØ-
Sg
log Pdata
2Ar
2
High Action
Action
I.IDD Samples
Transition Paths
Figure Proposed Onsager-Machlup Action Optimization Schematic. (Left) Atomistic generative models produce statistically independent samples via integration along a learned vector field, from which a score estimate, 8g- V Vlog/usuxx)) can be extracted. The score can be interpreted as drift term the stochastic dynamics induced by the generative model. (Right) This connection can be leveraged repurpose atomistic generative models to find high-pre peh-prrbbbbility transition paths between samples on the data manifold by minimizing the OM action functional, Eq. (14). The OM action has three terms which prioritize (I) low distance between adjacent points the discretized path, (II) low-norm drifts, and (III) convexity of the underlying energy landscape. The variables as follows: path point i (x()), trajectory length (L), timestep (At), friction ((), diffusion (D), latent time (T). The underlying score Sg- remains frozen throughout OM optimizaaion. Although valid for any data distribution, in the special case of Boltzmann Boltzmann-distrbbute distributed data, our approach has the natural interpretation of transition path sampling potential energy surface with atomistic force field.
from first principles by minimizing the OM action with gradient-based optimizers. In the specific case where the data Dolzzmann-distributed. the learned score approximates the underlying atomistic force field (Arts al., 2023), and our approach has the direct, physical interpretation of TPS on a potential energy surface at a finite temperature.
5. we build intuition show accurate estimation of reaction rates and committor functions using our sampled transition paths. On alanine dipeptide (Section 5.2), method recovers accurate free energy barriers and outperforms metadynamics and shooting algorithms in efficiency. For fast-folding proteins (Section 5.3), OM action minimization with diffusion or flow-matoe fow-matching models yields transition path ensembles closely aligned with reference MD at a fraction of the cost of traditional MD. Finally, we show that OM optimization with generative models trained tetrapeptide configurations enables zero-shot TPS new sequences (Section 5.4). Overall, our work demonstrates the promise of repurposing pre-trained generative models as a general-purposs. approach for transition path sampling.
Our approach has a number key advantages:
1. Scalability: We do not require any training procedure specific TPS, and instead exploit pre-trained gen- erative models. This makes our approach scalable models and datasets continue to grow, and generalizable to new systems without retraining.
2. Flexibility: We can incorporate physical parameters such as the time horizon and temperature when sampling paths without modifying the underlying generative model.
2. Related Work
Transition path sampling. Traditional approaches to TPS like umbrella sampling (Torrie & Valleau, 1977) and metady- namics (Laio & Parrinello, 2002a) employ biasing potentials along representative collective variables (CV). However, defining suitable CVs is challenging near transition states, even with automated approaches (Sultan & Pande, 2018;
3. Efficient Diversity:
We leverage the stochasticity of generative models to produce diverse candidate paths more efficiently than traditional methods.
We validate our approach on systems of increasing complex- ity. Starting with the 2D Müller-Brown potential (Section
2

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
Sípka et al.,, 2023). Meanwhile, shooting techniques (Mullen al., 2015; Borrero & Dellago, 2016; Jung al., 2017; Bolhuis& & Swenson, 2021), which use a Metropolis-Hastingg criterion to sample trajectories, suffer from slow sampling, high rejection rates, and the need for expensive simulations. While ML approaches, including reinforcement learning (Das et al., 2021; Rose et al., 2021; Singh & Limmer, 2023; Seong et al., 2024; Liang et al., 2023), differentiable simula- tions (Sipka al., 2023), and h-transform learning (Singh & Limmer, 2023; Du et al., 2024), have been used to design CVs or biasing potentials with promising results, they require ex- pensive sampling procedures, must be retrained for each new system of interest, and do not exploit atomistic simulation data. Minimizing the OM action has been used for TPS in low- dimensional systems (Faccioli et al., 2006; Vanden-Eijndee & Heymann, 2008; Autieri et al., 2009; Fujisaki al., 2010; Lee al., 2017), but has faced computational challenges, such as a lack of integration with modern auto-differentiation frameworks and the inability to perform gradient-based op- timization for larger systems (a Beccara et al., 2012). To our knowledge, our work is the first to propose a gradient-buut based optimization of the complete OM action, and to connect it to the stochastic dynamies induced by generative models. For more extensive review of TPS methods, see Appendix A.
with the intrinsic manifold structure of the data. Our OM optimization approach can be seen novel interpolation mechanism which leverages the inductive bias of stochastic dynamics to generate high-probabillty transition paths.
3. Theory
We now introduce the OM action as a way to compute path probabilities under particular stochastic differential equation (SDE), and describe its application to score-based generative modeling and transition path sampling.
3.1. Probability of paths under stochastic dynamics
We introduce the following constant variance SDE which will underpin our proposed framework for TPS:
x Φ (x)d
(1)
where (x)):R* IR* is a drift function, W, is a standard Weiner process, and D.5 > are scalar constants governing diffusion noise levels and damping respectively. We consider drifts which can be written as the gradient of a scalar: (xx) =
8x where o(x) 1²* - R. By solving a Fokker-Planck equation for the time-varying state distribution p(x.t), we can obtain the probability of path x(-) = {x(t)} sampled from the SDE in Eq. (1) (see
ad(x)
Atomistic generative models. Generative models can produce unbiased, independent samples from the config- urational ensemble of molecular systems, pioneered by Boltzmann generators (Noé et al., 2019) and since further developed for proteins (Arts et al., 2023; Zheng et al., 2024:
Appendix C for complete details):
((x(())xee
(2)
To maximize this probability with respect to a path, we can equivalently minimize the negative log probability S, which is called the Onsager-Machlup action functional. This is the stochastic analogue of the well-known principle of least action from optics and quantum mechanics (Rojo et al., 2018). Since only consider discretized paths in this work, we focus the discretized form of the OM action:
Jing et al., 2024a; Lewis et al., 2024; Schreiner et al., 2023),
small molecules (Huang et al., 2024; Schneuing et al., 2024;
Reidenbach & Krishnapriyan, 2024; Igashov et al., 2024), and materials (Zeni et al., 2023; Zheng et al., 2024; Xie et al., 2022). Generative models are typically trained to match the distribution of atomistic configurations from large-scale datasets, including structural databases (Bank, 1971) and long-timescale MD simulations (Lindorff-Larsen et al., 2011; Vander Meersche et al., 2024). Recent works adapt generative models to produce more diverse samples (Corso et al., 2023), perform rare event sampling (Falkner et al., 2023), perform MD simulations using the connection between diffusion mod- els and force fields (Arts et al., 2023), and learn generative models directly over trajectories (Jing et al., 2024b).
Definition 3.1. For a discrete path X (x."))....x( generated by the SDE in Eq. (1) with drift & and timestep size At, the discretized form of the Onsager-Machlup action functional is given by:
S[X]
A[X]+ + -
2D
Interpolations in generative models. Analagous to TPS, interpolation has been used to evaluate the smoothness and continuity of learned data manifolds and to generate realistic transitions between data points using generative models. While linear interpolation in model latent spaces is known to capture some continuity (K (Kingma & Welling, 2013; Good- fellow et al., 2014), geometric techniques such as geodesic interpolation and optimal transport (Arjovsky et al., 2017: Ar- vanitidis et al., 2018; Lešniak et al., 2018; Michelis & Becker, 2021; Struski et al., 2023; Psenka et al., 2024) better align
[ 1
x (1+1)
x (6)
2At
(3)
A
B 20²
Σ
$
DAI
C[X] =
Σ
V x (x (1)
5 1=1
3

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
The three summands of the OM action each have intuitive interpretation. Term A encourages smooth transitions along the discretized path. Term B encourages paths remaining in regions with low-norm drifts, which are equilibria or saddle points of the underlying dynamics. Finally, term C encourages paths to remain in regions with low divergence of the drift, which can be interpreted as entropically favoring regions of convexity in the landscape of Φ, where dynamics more stable. The parameters Ç. At, and D control the relative contribution of these three terms in a physically intuitive manner. For instance, at larger values of At, the contribution of term A is diminished, consistent with the intuition that larger "jumps" are more probable with larger timesteps. the limiting case of negligible diffusivity D (analogous to temperature, see Appendix C), the divergence term can be omitted, yielding the Truncated OM Action:
matching. Specifically, we show that these models induce set of stochastic dynamics whose drift is given by their learned score function. This provides powerful framework to reason about probabbbbbyylity transition paths.
3.3.1. STOCHASTIC DYNAMICS
UNDER DENOISING DIFFUSION MODELS.
Denoising diffussio probabilistic models (DDPM) (Ho et al., 2020) are a class of score-based generative models that learn how to de-noise corrupted samples. The DDPM objective Eq. (25) is closely linked to the score-matching objective (Vincent, 2011; Song & Ermon, 2019) for training a score model 80 80 (x,7):
IR& x IR+ - RR* parameterized by 0:
(Sm xmppp [|so(x,))-Vloge))) (x))1
(7)
Definition 3.2. The truncated OM action of a discretized path X is given by:
where Vlogp, is the score of the noised marginal distri- bution pr time T. This establishes the connection between the score and the optimal noise model with parameters 8*: co- (x. x - (x). See Appendix B.2 for detailed statements and proofs. In order to use the OM action to compute path probabilities with a DDPM, we must construct a surrogate SDE in the form of Eq. (I), such that paths under this SDE have high likelihood under the data distribution used to train the model. While the denoising (i.e., sampling) process of a DDPM (see Appendix B.1) may appear to be natural candidate, closer inspection reveals that it is un- suitable, as it optimizes for different likelihoods at different points of the trajectory. A large portion of the denoising trajectory thus has low likelihood under the data distribution. Therefore, we need to consider an alternative approach.
1 (A[X B((I). 2D
S[X]
In this work, we use both the truncated and full OM actions depending on the particular system considered.
3.2. Transition path sampling in molecular systems
The connection between the SDE in Eq. (1) and TPS is straightforward. Formally, in TPS we consider d- dimensional molecular systems with Np-many particles inter- acting under potential energy function U (x):: (x)::R%) R xd + R, where x E Q is a configuration of the system, and Ω R, xd is the configuration space. The goal of TPS is to, for the given system and temperature, find most likely paths (x()) 1 J Josist over time horizon Tp, where At is the simulation timestep and L= Te is the number of discretization points in the trajec- AF tory. The trajectory traverses the two endpoints x0)) E AACS Ω, x( S2, where A.B typically represent distinct minima on the potential energy surface Uxx). The underlying parti- cle dynamics are governed by the SDE in Eq. (1), known in this context as overdamped Langevin dynamics. The scalar o and gradient-brart based drift &b terms have a clear physical in- terpretation as the potential energy and forces, respectively:
Iterative denoising and noising as candidate SDE. Another hypothesis for constructing SD to leverage the process of iterative one-step denoising and noising at a fixed time marginal T of the diffusion process. Intuitively, this bal- ances the sewiihooddmmiimizzzig drift of the denoising step with the stochasticity of the noising step. Specifically, we consider the following iterated denoise-noiss updates, where eg (x,7) is the trained denoising model from the DDPM:
x (i.mid))
1
x
B- ar
co cx(x(()) .7)
+ +V VBz2
V1
Br
VT-ar
(8)
((x) : Ux),
(5)
x 1) = v1 11-111111 B (f.mid) + +V VB
(xx) = F (x).
(6)
(9)
where 2,2' (O,(),) a denote the usual diffusion model noise schedule variables, and ā II az. Combining the two updates yields a single update equivalent distribution, yields:
If the forces Fxx) are known, the OM action (Eq. (3)) can be used to compute the probability of paths connecting endpoints ( x (1) under the governing dynamics.
writing so (X,T)= = Tā yicd)s
3.3. Score-based generative modeling
(++1) (x+1)=((0).._ B, I ar-l so (x() 5)+ 123, -32
We now describe the connection between the SDE in Eq. (1) and generative models, namely denoising diffusion and flow
v a-
(10)
4

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
where = ~ (O,,)) Taking the continuum limit of DDPM sampling discretization steps infinity, we get that -0,, and noting that 23, - _ x 23, and är -, at this limit, we see that Eq. (10) is an Euler-Man Haler-Maruyama discretization, with timestep 3r, of the following SDE:
og = =1 00 01 0 1, and 0,,, -o are both strictly increasing functions.
While extracting the learned score Sp from DDPM is straightforward via sg - - -L/VT-ar--)))(x it is less clear how to do so for flow matching. However, we note the following:
+ W.W.
(11)
Note that the above construction holds for any (Ta where Ta is the maximum diffusion time. A similar derivation can be found in Arts et al. (2023) T=0.
By Eq. (7), the targets for the denoising model e (x.7) DDPM are equivalently the negative scores of the noised distribution, Vlogp-(x))
Eq. (11) is equivalent (up to constants) to Eq.(1):
The targets ux(x) for the flow matching model (x().T) can be converted to scores of the flow marginal distribution Vlogp- (x) through the following formula (see Appendix B.3 for the proof):
o( xx) lolopp,(x) (x). Φ xx x (x,T) Vlogp,
(12)
(13)
This means that, as introduced in Eq. (2) and Eq. (3), the likelihood of discrete paths = (x(*) OSAST. under the constructed SDE Eq. (11) can be evaluated via the OM action defined by the trained DDPM score Sø- (x,T):
Vlogp-(x)= 0,0,a,-a,e x-x-uulu ur ux(xx) (xx)
at
( (i x
(17)
We can thus extract an approximate score sig a trained flow matching model ug. (x,T) via,
PM (x,T) from
S(X;0)== 2D
x ((+1)
x00
2At
(14)
år xo (x,7). xxxx))((x)) (18)
+
58- (x()),,7 2(²
At
.r)
||?
D&t
V so. (x(())
FM (X,T) 0,0+0, à.07 at
By inserting spM (x, T) into the denoise-noise process defined in Eq. (11), we again obtain an SDE of the form of Eq. (1). Hence, we can use the OM action to compute the log-probabilities of paths between arbitrary datapoints.
where 1, and At Br. Note that B, can be effectively tuned as hyperparameter by changing the sampling discretization fidelity. In our specific scenario of TPS, we can set these parameters to physically interpretable values, with Al, C. and D corresponding to the MD simulation timestep, friction coefficient, and diffusion coefficient, respectively (see Section 3.3.3).
3.3.3. PHYSICAL INTERPRETATION FOR BOLTZMANN-DISTRIBEE DATA
By assuming underlying dynamics of the form Eq. (11), our framework combining generative models and the OM action can be used to compute the log pog--obbbbiliiees of paths between samples from any data distribution with a well-defined score. However, in the special case where the data used to train the generative model are Boltzmann distributed, i.e., p(x) x exp knT ), where kB is Boltzmann's 's constant and '' is the temperature, the leameed score Sø- (x,7 =0) : N, R Ny is interpretable as physical, atomistic force field:
3.3.2. EXTENSION TO FLOW MATCHING
We now link OM cattionmiiiiizztioon with broader class of generative models beyond DDPM. Flow matching models (Lipman et al., 2022) are a natural choice due to their strong performance in generative modeling tasks across modalities (Jing al., 2024a; Polyak et al., 2024). Similarly to DDPM, flow matching generates samples through repeated integration process over a learned vector field. For affine flows considered in this work, the training objective for the learned velocity field us (I,T) takes the form,
[(xx),
Sg (x,7 0)) vo Vlogp(x) x VUx))=
(19)
The constructed dynamics in Eq. (11) at T reduce to overdamped Langevin dynamics governing particles po- tential energy surface. Thus, can directly set the constants in the OM action (At, Ç, D, T) to physical values used in MD simulations for a given atomistic system. Under our OM framework, finding bigh-probability paths between points on Boltzmann dalizmann-distributeedd data manifold thus directly aligns with the conventional notion of TPS for atomistic systems.
[|uo(x;))) u, (x)
(15)
u (xx)= E. [d,xi +ox00 x= xx
NorPPa Xi
(16)
where Po and P1 are the source and target distributions,
0r, 0r [0, 1] - [0, 1] define a curve from Xo to Xi:
5

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
4. Repurposing A Atomistic Generative Models via Action Minimization
Algorithm I OOnsager-Machlup Transition Path Optimization with Generative Models
Input:
We now introduce our OM optimization approach to produce high probability transition paths between atomistic structures using pre-trained generative models. Given a pre-trained generative model with a fixed score function Sø-, and two atomistic configurations, x x("), xx(2) E IR xd where x E A.x((*) e B, we aim to sample transition path x= [xx*) (x(), ]i consisting of T. discrete steps. Our core inductive bias is to interpret candidate paths as realizations of the denoise-noise SDE in Eq. (11), enabling tractable computation and optimization path log-likelihoods via the OM action. Our approach proceeds in three primary steps: 1) computing initial guess path, 2) performing OM optimization, and (in some cases) 3) decoding the optimized path back to the configurational space Ω. Algorithm summarizes the procedure without the decoding step.
Optimization time Topt
3: Generative model with time-conditional score sp. (:/Tept)
score
4: Two atomistic configurations x(") A C R RNxdd
x!D) EB RBRR
x
5: Compute initial guess X = -x(00) (x(0),,,,x00)) =
InitialGuess( (0),, (Tmina)))))))) (Algorithm 2)
6: while not converged do
7: Compute the OM action, S[X;0*), with the learned vector field Sg- (--Tope). ). using Eq. (14).
Update X < optimizer(X (X, xx={(x;01)) (keeping
the endpoints x0) (x)) fixed)
end while
50
LLD. Samples
Computing an initial guess path. The choice of initial path connecting the endpoints ( E A and (44 B cru- cial in determining the quality of the subsequently optimized path. Naïve linear interpolations in configurational space likely unphysical, since plausible atomistic structures typically lie on a highly non-convex, low-dimensional manifold of Ω. Instead, we opt to linearly interpolate at a latent level Timitial of our pre-trained generative model, which known to produce samples closer to the data manifold (Ho et al., 2020). After interpolation, we can either (1) decode the interpolated path back into configurational space and subsequently optimize the OM action there, (2) directly optimize the OM action at the latent level Tinitial Both are justified, since the proposed dynamics in Eq. (11) are valid for any r. See Appendix G for complete details.
OM Optimization Diffusivity-0
40
OM Optimization Diffusiivty=1100
30
OM Optimization Diffusivity=4..0
20
10
90
20
40
50
X
Figure 2. OM optimization with diffusion model the 2D Müller-Brown potential. Individual points along the OM-optimized paths are shown as dots. Increasing the diffusivity D causes the path to cross at a higher barrier. An equivalent number of L.L. samples (red) fails to sample the transition region.
Optimization of the OM action. Starting from the initial guess, we find paths which have high probability (low action) under the SDE in Eq. (1I) inddcced by the generative model. The optimization problem simply becomes.
X =
argmin S[X;0")
(20)
Decoding back to configurational space. If OM- optimization was performed at non-zero latent time Tinidals we decode the final path, obtained after K iterations of gradi- ent descent, back to the configurational space. If Ifoptimization was performed in configurational space (i.e., decoding was already done in the first step), then this step is skipped.
X= X= (x(0)) ( 1ejo.41
where S[::6" 0*] is the generative model action in Eq. (14) (computed with the frozen, pre-trained score sg-) ), and (0) and x() are kept fixed. We approximate this minimum via gradient descent on the path until the action converged. Since S[=:8*] a discretized integral, the entire trajectory optimized in parallel, which is amenable mevti-device acceleration. In line with previous work (Arts et al., 2023), we find that a small, nonzero value of Topt leads to better alignment with the true forces, and thus treat it as hyper- parameter (see Appendix G). To accelerate computation of the divergence term in the OM action, we use the Hutchinson estimator (Hutchinson, 1989) (see Appendix G for details).
5. Results
We now present the results of our OM optimization approach for TPS with pre-trained generative models. In all cases other than the Müller-Brown potential (Section 5.1), we perform OM optimization in configuration space, so we skip the final decoding step and use physical parameter values. In Appendix F, we demonstrate use case beyond generative modeling, namely using a classical force field to generate
6

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
(a)
Sampled Transition Paths
(b)
Potentaal of (kcalmol)
Energy Profile
energy minima, and the path crosses a higher energy barrier, consistent with the larger scale of thermal fluctuations increased temperatures. Meanwhile, 2500 i.idd samples (shown in red) from the diffusion model sample only the three energy minima, and fail to sample the transition region. This provides a proof-of-concepp that our OM optimization procedure can be used to repurpose generative models to sample transition paths without specialized training. See Appendix H for additional results and analysis on the MB potential, including with flow matching model.
w s uP
N
101
3
2
0 PCA CV (PC1)
2
9 (Phi)
Figure 3. OM optimization and free energy barrier estimation on alanine dipeptide using a pre-trained diffusion model. (a) Sampled alanine dipeptide transition paths from OM optimization with pretrained diffusion model, overlaid on a Ramachandran plot of sample density at 600K. (b) Free energy barrier for selected path, estimated via umbrella sampling. is approximately kcal/mol, in line with our metadynamics calculations (I).
Committor and Transition Rate Estimation. Committor functions and rates are fundamental quantities in the study of transition events (Vanden- Eijnden et al., 2006). The committor function (xx) is defined as the probability that trajectory initiated at x = X reaches the target state B before the starting state A. Transition paths obtained via OM optimization can be used as an enhanced sampling method to accurately compute the committor function, and subse- quently the transition rates, on the Müller- potential (see Appendix H for complete details). Specifically, we initiate MD simulations from points along the OM-optimized paths shown in Fig. 2 to collect a dataset of samples Dusin. We then train a committor function qa (x) by solving a functional optimization problem given by the Backward Kolmogorov Equation (BKE) over the sampled points, and finally compute an estimate of the reaction rate using the trained model via the relation
Table Benchmarking the speed of OM optimization against traditional enhanced sampling baselines on alanine dipeptide. We report the number of force field evaluations required per 1,000 step transition path (for OM optimization
instead report the number of score function ovaluations). OM optimization requires significantly fewer evaluations and thus lower computational cost per path than the traditional approaches.
Method
Puth (1)
Runtime/ Path (L)
MCMC (Two-Way Shooting)
No
2: IB
2 100 hours
Melady
Yes
IM
hours
Opt. (Diffusion Modell (ours)
et al., 2006):
KBT γ
Kg
(x)x)(x))
=xx V Vx2o(x)) (x)
(21)
all-atom transition paths for Chignolin and Trp-Cage.
5.1. 2D Miller-Brown potential
where (-)a denotes an ensemble average over the configura-
We first demonstrate our method on the 2D Küiler-Broown (MB) potential (Müller & Brown, 1979), a canonical test system for TPS with global minimum and two local minima, separated by saddle points.
tional space It and E denotes an empirical mean. Using this procedure, we obtain a transition rate estimate of 1.3 x 10-5, compared with the true rate of 5.4 x 10-5, It often challeng- ing to compute the reaction rate even within the correct order of magnitude (Rotskoff et al., 2022; Hasyim et al., 2022), sug- gesting the promise of our OM optimization method to enable accurate rate estimation more challenging systems.
Problem setup. Using a denoising diffusion model pretrained on samples from the underlying potential energy surface, we generate transitions between the two deepest energy minima using OM optimization. We generate the initial guess path via linear interpolation at Tinitial 8 (the maximum diffusion model time is Ta = 100).. We perform 200 steps of OM optimization directly at Timiial= = using Tope = 8, and finally decode the path back to the data distribution via the denoising process.
5.2. Alanine dipeptide
Alanine dipeptide is a classic benchmark system for TPS, with 22 atoms and CVs described by the dihedral angles d.
Problem setup. We start with a denoising diffusion model pre-trained on samples from the alanine dipeptide potential energy surface. Using this model, we perform OM optimization to find transition paths between the two standard minima defined in the CV space.
Results. As shown in Fig. 2. OM optimization with the truncated action yields a transition path (shown in blue) between the energy minima which passes through the lowest energy barrier. Due to the stochastic decoding process, the samples around the transition path exhibit natural diversity. Increasing the diffusion coefficient D = 7M
Results. We successfully sample two likely transition paths between the metastable basins (Fig. 3a). In Table 1, we benchmark the computational efficiency of OM optimization against traditional enhanced sampling techniques, namely metadynamics (Laio & Gervasio, 2008) and MCMC-
kn? results in
qualitatively different transition paths (shown in orange and yellow). Samples are more concentrated in the three
7

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
(a)
Reference Energy Surface
Sampled
(b)
(d)
Transition
Transition State Ensemble
Markov State Model Metrics
BBA
0.5
Negativll Logg Liaet
JSD & Volld Path
0.6
Commitoo Probbbiity
Free energy kaf
04
Villin
(c)
Sampling Time
0.4
Will Clock Time (CONH-our)
0.02
Trp- Cage
30n
0.0
Kreas
WAe TAMSOOS
Negathe Lnailtoos
Antrage
0.00
Dierge
Interpolation
Diffusion Model
75ys MD
10ps MD
Interpolation
50us
Sjss
Sampled Transition Paths Transition State Ensemble
Matching
25us MD
MD
101
100us MID
Figure OM optimization with diffusion and flow matching models trained on coarse-grained, fast-folding proteins. (a) Reference free energy surfaces of fast-folding proteins, alongside transition paths produced by OM optimization (yellow) overlaid against the landscape the empirically computed committor function xxx). The transition state ensemble (black) the set {x: [x:0.45 0.45 S g(x) S 0.55). (b) Samples from the predicted transition state ensemble of BBA. (c) Runtime of OM optimization and varying lengths of unbiased MD simulation of BBA. MD simulation is performed with diffusion model score an approximate force field, as in Arts et al. (2023). (d) MSM results averaged across 5 fast-folding proteins. Comparisons
with respect reference, unbiased MD simulation of varying lengths. Plotted are Jensen-Shannon Divergence (a measure distributional dissimilarity) between the sampled and reference path probability under the reference MSM). and average
MSM state distributions, fraction of sampled paths which are valid (i.e., have transition negative log likelihood under the reference MSM (indicating the realism the paths).
two-way shooting (Bolhuis & Swenson, 2021). and find that OM optimization considerably more efficient. As shown in Fig. 3b, we can also use the transition paths resulting from OM optimization as a natural CV for umbrella sampling (Torrie & Valleau, 1977), from which we can accurately and efficiently estimate free energy profiles along the transition path. See Appendix for complete details on model training, OM optimization, and free energy calculations.
1. Transition negative log likelihood. The negative log likelihood of MSM state transitions under the reference MSM, averaged over all paths with non-zero probability.
2. Fraction of valid paths. The fraction of paths with non-zero probability under the reference MSM.
Jensen-Shannon divergence. The JSD (distributional dissimilarity) between the distribution of states visited by the generated paths and those sampled from the reference MSM.
5.3. Fast-folding coarse-grained proteins
We next consider proteins exhibiting fast dynamical transitions, for which millisecond reference MD simulations were performed in Lindorff-Larsen et al. (2011).
We compute these metrics for our generated transition paths, as well as for 1 100pss subsets of the reference MD simula- tions To compare wall-clock time for generating transition paths, since reference simulations were performed all-atom resolution and their speeds unavailable, we follow Arts et al. (2023) and run coarse gaasse-griindd MD using the generative model's learned score function as an approximate force field. Results. As shown in Fig. 4a, OM optimization yields diverse transition paths which intuitively pass through high density regions of the free energy landscape, projected onto the two slowest Time Independent Component (TIC)
Problem setup. We adopt coarse- coarse-graining graining (CG) scheme which represents each amino acid with the position of its C. atom. We utilize the pre-trained diffusion models from Arts et al. (2023), and we train our own flow matching models. Separate models are trained for each protein. To facilitate analysis and interpretation of results, we divide the conformational space into discrete states and make use of Markov State Models (MSMs) (Prinz et al., 2011; Noć et al., 2013) to obtain state transition probabilities. Similar to Jing et al. (2024b), we evaluate the quality transition paths by discretizing them over the MSM states and computing the following metrics (see Appendix for complete details):
We did not consider enhanced sampling baselines for the fast-folding proteins due to their reliance energy function, which generally not available for our level coarse-graining.
8

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
paths using this data-starved model (see Appendix J for more details); suggesting that our approach can be useful even if the underlying data distribution is not an exhaustively sampled Boltzmann distribution.
1.0
3.0
JSD & Valid Rate
0.8
2.5
Negative Log Likellhood
2.0
5.4. Generalization to new tetrapeptides
0.6
As a final evaluation, we consider all-atom tetrapeptide sys- tems, which exhibit interesting dynamics and pose the chal- lenge of generalization to held-out amino acid sequences.
0.4
1.0
0.2
0.5
Problem setup. We train a flow matching generative model on approximately 3,000 tetrapeptides simulated in Jing et al. (2024b), and apply our OM optimization procedure to generate an ensemble of 16 transition paths for each of 100 tetrapeptides not seen during training. We use the same MSM-based metrics as in Section 5.3 to evaluate the quality of generated transition paths.
0.0
Jensen
Valid Transition
Average Transition
0.0
Shannon
Path Rate (1)
Negative Log Likellhood (4)
Divergence (4)
OM Interpolation - Flow Matching Model
30ns MD
5ns MD
100ns MD
20ns
2ns
50ns MD
10ns MD
5. OM optimization on unseen tetrapeptide sequences. OM optimization with a flow matching model yiclds transition paths which compare strongly with variable-length MD simulations, indicating the generalization potential of our approach. Plotted are the Jensen Shannon Divergence (a measure distributional dissimilarity) between state distributions visited by the reference and generated paths, the fraction of valid (non-zero probability) paths under the reference MSM, and the negative log likelihood of state tramsitions under the reference MSM (indicating path realism).
Results. As shown in Fig. 5, OM optimization achieves MSM metries which are competitive with MD simulations of 50-1000 ns, which are considerably more computationally expensive to generate. This suggests the promise of OM optimization to generate transition paths on atomistic systems not explicitly seen during training. See Appendix K for further details and path visualizations.
6. Conclusion
(Noé et al., 2013) axes. We can also vary the physical parameters of the OM action (e.g., the time horizon Tpp) to obtain paths traversing different regions of phase space (see Appendix J). For the BBA protein, the paths robustly sample the transition state ensemble, empirically defined by the level set 0.45 0.55) of the committor function (Fig. 4b). Sampling transition paths of the BBA protein with OM optimization requires considerably less wall-clock time than using the diffusion model's 's learned score as a coarse-grained force field and performing unbiased MD simulations (Fig. 4c). Across all proteins and both classes of generative model (diffusion and flow matching), OM optimization yields a higher percentage of valid paths and lower transition negative log likelihood under the reference MSM, compared with unbiased, reference MD simulations of any of the considered lengths up to 100 jas (Fig. 4d). The JSD is also lower than any MD simulation length up to 50 jas, indicating that the sampled paths traverse similar distribution of MSM states as the reference simulations. See Appendix J for comparisons to transition trajectories found in the reference, unbiased MD simulations.
We have presented a method to repurpose atomistic genera- tive models for transition path sampling by finding paths over the data manifold which minimize the Onsager-Machlup action under the model's 's learned score function. Our ap- proach, which requires no TPS-specific training procedure, aligns well with the growing trend of leveraging large-scale, well-tested, general-purpose generative models direction already standard in the language and vision communities.
Limitations. Our approach does not provably sample the full posterior distribution over paths, as in traditional shooting methods and recent ML. approaches (Du et al., 2024). However, we sample diverse paths by exploiting the stochastic generative model encoding and decoding process. Initiating traditional MD or umbrella sampling simulations is another way to explore the potential energy surface around the OM-optimized paths (see Appendix I).
Future work. Incorporating OM optimization into larger generative models trained on more diverse data (Lewis et al., 2024: Jing et al., 2024a) is a natural area for future develop- ment. Given the success of large-scale, co-evolutionary mod- eling of proteins (Jumper et al., 2021), it would be interesting to investigate the extent to which pre-training generative models on large structural databases enables TPS on unseen systems. More broadly, OM miniiinnmiiiiiziiionnn could be a powerful framework to generate interpolation paths in a va- riety of data modalities, including images, videos, and audio.
Robustness to sparse data in transition regions. To simulate the scenario in which transition states are not well-represented in the training data, we retrain diffusion models on datasets from which 99% of configurations with committor probability between 0.1 and 0.9 are removed. OM optimization is still able to sample plausible transition

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Impact Statement
Autieri, E., Faccioli, P., Sega, M., Pederiva, F,, and Orland,
H. Dominant reaction pathways in high-dimensional systems. The Journal of chemical physics, 130(6), 2009.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Bank, P. D. Protein data bank. Nature New Biol, 233(223): 10-1038, 1971.
Acknowledgments
Baratta, A., Dean, J. P.,, Dokken, J. S., Habera, M., HALE, J., Richardson, C. N., Rognes, M. E., Scroggs, M. W., Sime, N., and Wells, G. N. Dolfinx: the next generation fenics problem solving environment. 2023.
The authors thank Rasmus Lindrup. Aditya Singh, Muham mad Hasyim, Kranthi Mandadapu, Simon Olsson, Soojung Yang, Lukáš Grajeiar, Johannes Dietschreit and David Lim- mer for helpful discussions that benefited this paper, well as Bowen Jing for assistance with reproducing the min-flur endpoints for the tetrapeptide evaluations. The authors acknowledge support of this work from the U.S. Department of Energy, Office of Science, Energy Earthshot initiatives as part of the Center for lonomer-based Water Electrolysis at Lawrence Berkeley National Laboratory under Award Num- ber DE-AC02-0CCHI 1231. This work was also supported by the Toyota Research Institute as part of the Synthesis Advanced Research Challenge. MP and MŠ were supported by the Czech Science Foundation, project 23-05736S.
Bolhuis, P. G. and Swenson, D. W. Transition path sampling markov chain monte carlo of trajectories: Recent algorithms, software, applications, and future outlook. Advanced Theory and Simulations, 4(4):2000237, 2021.
Bolhuis, P. G., Chandler, D., Dellago, C., and Geissler, P. L. Transition path sampling: Throwing ropes over rough mountain passes, in the dark. Annual review of physical chemistry, 53(1):291-318, 2002.
Borrero, E. E. and Dellago, C. Avoiding traps in trajectory space: Metadynamics enhanced transition path sampling. The European Physical Journal Special Topics, 225: 1609-1620, 2016.
References
Beccara, S., Škrbić, T., Covino, R., and Faccioli, P. Domi- nant folding pathways of a ww domain. Proceedings of the National Academy of Sciences, 109(7):2330-2335,, 2012.
Chen, W. and Ferguson, L. Molecular enhanced sampling with autoencoders: On-the-fly collective variable discovery and accelerated free energy landscape exploration. Journal of computational chemistry, 39(25): 2018.
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J.,
Bambrick, J., et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630 (8016):493-500, 2024.
Corso, G., Xu. Y., De Bortoli, V., Barzilay, R., and Jaakkola, T. Particle guidance: non-iid diverse sampling with diffusion models. arXiv preprint arXiv:2310.13300 2023.
Adib, A. B. Stochastic actions for diffusive dynamies: Reweighting, sampling, and minimization. The Journal of Physical Chemistry B, 112(19):5100-591 05 2008. ISSN 1520-6106. doi: 10.1021/jp0751458.
Cugliandolo, L. and Lecomte, V. Rules of calculus in the path integral representation of white noise langevin equa- tions: the onsager- machlup approach. Journal of Physics A: Mathematical and Theoretical, 50(34)3345001, 2017.
Arjovsky, M., Chintala, S., and Bottou, Wasserstein gan. In Proceedings of the International Conference on Machine Learning. pp. 214-223, 2017.
Darve, E. and Pohorille, A. Calculating free energies using average force. 2001. doi: 1..1063/..1419999
Das, A., Rose, D. C., Garrahan, J. P., and Limmer, D. T. Reinforcement learning of rare diffusive dynamics. The Journal of Chemical Physics, 155(13), 2021.
Arts, M., Garcia Satorras, V., Huang, C.-W., Zugner, D.,, Fed- erici, M., Clementi, C., Noé, F., Pinsler, R., and van den Berg, Two for one: Diffusion models and force fields for -aarse-grained molecular dynamics. Journal of Chemical Theory and Computation, 19(1 19(18)66151-615) 2023.
Dellago, C., Bolhuis, P. G., Csajka, F. S., and Chandler, D. Transition path sampling and the calculation of rate constants. The Journal of Chemical Physics, 108(5):1964- 1977, 02 1998. ISSN 0021 19606. doi: N0.1063//..47550
Arvanitidis, G., Hansen, L. K., and Hauberg, S. Latent space oddity: On the curvature of deep generative models. In Proceedings of the International Conference on Learning Representations, 2018
Dellago, C., Bolhuis, P. G., and Geissler, P. L. Transition path sampling. Advances chemical physics, 123:1-78 2002.
10

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
Doerr, S., Majewsk, M., Pérez, A., Krämer, A., Clementi, C., Noe, F., Giorgino, T., and Fabritiis, G. D. Torchmd: A deep learning framework for molecular simulations, 2020.
Holdijk, L., Du, Y., Hooft, F,, Jaini, P., Ensing, B., and Welling, M. Stochastic optimal control for collective vari- able free sampling of molecular transition paths. Advances in Neural Information Processing Systems, 36, 2024.
Du, Y., Plainer, M., Brekelmans, R., Duan, C., Noé, F., Gomes, C. P., Aspuru-Guzik, A., and Neklyudov, K. Doob's lagrangian: A sample-efficiett variational approach to transition path sampling. arXiv preprint arXiv:2410.07974.20000 2024.
Huang, L., Xu, T., Yu, Y., Zhao, P., Chen, X., Han, J., Xie, Z., Li, H., Zhong, W., Wong, K.-C,, et al. A dual diffusion model enables 3d molecule generation and lead optimization based on target pockets. Nature Communications, 15(1):2657, 2024.
Faccioli, P., Sega, M., Pederiva, F., and Orland, H. Dominant pathways in protein folding. Physical review letters, 97 (10):108101, 2006.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Com- munications in Statistics Staristies-Simulation and Computation, 18()):059-1076, 1989.
Falkner, S., Coretti, A., Romano, S., Geissler, P. L.., and Dellago, C. Conditioning boltzmann generators for rare event sampling. Machine Learning: Science and Technology, 433):035050, 2023.
Igashov, I., Stärk, H., Vignac, C.,, Schneuing, A., Satorras, V. G., Frossard, P., Welling, M., Bronstein, M., and Correia, B. Equivariant 3d-conditional diffusion model for molecular linker design. Nature Machine Intelligence, pp. 1-11,2024
Frenkel, D. and Smit, B. Understanding molecular simu- lation: from algorithms applications. Elsevier, 2023.
Fujisaki, H., Shiga, M., and Kidera, A. Ensager-Machile action-basee path sampling and its combination with replica exchange for diffusive and multiple pathways. The Journal of Chemical Physics, 132(13): 134101, 04 2010.
Jing, B., Berger, B., and Jaakkola, T. Alphafold meets flow matching for generating protein ensembles, 2024a.
Jing, B., Stärk, H., Jaakkola, T., and Berger, B. Generative modeling of molecular dynamies trajectories. arXiv preprint arXiv:2409.
ISSN 0021-9606. doi: /0.1063/1.3322002
2024b.
Gel' fand, I. M. and Yaglom, A. M. Integration in functional spaces and its applications in quantum physics. Journal of Mathematical Physics, 111):4869, 01 1960. ISSN 0022-2488, doi: 10. 17.1663.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zidek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):83-588) 2021.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Advances in Neural Information Processing Systems, pp. 2672-2680, 2014.
Jung, H., Okazaki, K.-i., and Hummer, G. Transition path sampling of rare events by shooting from the top. The Journal of chemical physics, 147(15), 2017.
Hasyim, M. R., Batton, C. H., and Mandadapu, K. Supervised learning and the fimite-temperature string method for computing committor functions and reaction rates. The Journal of chemical physics. 157(18), 2022.
Kabsch, W. Automatic processing of rotation diffraction data from crystals of initially unknown symmetry and cell constants. Applied Crystallography, (6(6))795-800 1993.
Kappler. J. and Adhikari, R. Stochastic action for tubes: Connecting path probabilities to measure- ment. Phys. Rev. Res., 2:023407, Jun 2020. doi: 0011039PhyeeeRRRecchc2.33333
Hendrycks, D. and Gimpel, K. linear units (gelus), 2023. //arxiv. org/////aooo.. aaaaabbb//666 08415.
Gaussian
URL https:
Henkelman, G.. Uberuaga, B. P., and Jónsson, H. A climbing image nudged elastic band method for finding saddle points and minimum energy paths. The Journal of Chemical Physics, 113(22)990119990) 12 2000. ISSN 0021-9606. doi: 10.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization, 2017.
Kingma, D. P. and Welling, M. Auto-encoding variational
bayes. In Proceedings of the International Conference on Learning Representations, 2013.
Ho, J., Jain, A., and Abbeel, Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851 2020.
Klein, L. and Noé, Transferable boltzmann generators. arXiv preprint arXiv:2406. /4426, 2024.

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machiin Functional
Kumar, S., Rosenberg, J. M., Bouzida, D., Swendsen, R. H., and Kollman, P. A. The weighted histogram analysis method for free-energy calculations on biomolecules. i. the method. Journal of computational chemistry, 13(8):
parameters from ff99sb. Journal of Chemical The- ory and Computation, 11())3696-3713, 2015. doi: 10. 10.1021/acsject.5500002.
Mauri, R. Non-Equilibrium Thermodynamics in Multiphase Flows. Soft and Biological Matter. Springer Netherlands, 2012. ISBN 9789400754614, URL https: /books. google cz bookss/ooos/// ? a22.a.
1011-1021, 1992.
Laio, A. and Gervasio, F. L. Metadynamics: method to simulate rare events and reconstruct the free energy in biophysics, chemistry and material science. Reports on Progress in Physics, 71(12):126011,2 2008.
Michelis, M. Y. and Becker, Q. On linear interpolation in the latent space of deep generative models. arXiv preprint arXiv: 0tti:2205.0088 03.63333, 2021.
Laio, A. and Parrinello, M. Escaping free-energy minima. Proceedings of the National Academy of Sciences of the United States of America, 99(20); 10 2002a.
Mullen, R. G., Shea, J.-E., and Peters, B. Easy transition path sampling methods: Flexible-length aimless shooting and permutation shooting. Journal of chemical theory and computation, 1166)24411244 2015.
Laio, A. and Parrinello, M. Escaping free-energy minima. Proceedings of the national academy of sciences, 99(20): 12562-12566, 2002b.
Müller, K. and Brown, L. D. Location of saddle points and minimum energy paths by a constrained simplex optimiza- tion procedure. Theoretica chimica acta, 53:75-93, 1979.
Lee, J. Lee, I.-H., Joung, I., Lee, J., and Brooks, B. R. Find- ing multiple reaction pathways via global optimization of action. Nature Communications, 8(1)315443, 2017. ISSN 2041-1723. doi: 10. 10.1008/ncomms115m
Noé, F., Wu, H., Prinz, J.-H., and Plattner, N. Projected and hidden markov models for calculating kinetics and metastable states of complex molecules. The Journal of chemical physics, 139(18), 2013.
Leśniak, D., Sieradzki, I., and Podolak, I. Distribution- interpolation trade off in generative models. In Interna- tional Conference on Learning Representations, 2018.
Noé, F., Olsson. S., Köhler, J., and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many- systems with deep learning. Science, 365(6457)xeaaw) 147, 2019.
Lewis, S., Hempel, T., Jiménez I Luna, J., Gastegger, M., Xie, Y., Foong, A. Y., García Satorras, V., Abdin, O., Veeling, S., Zaporozhets, I., et al. Scalable emulation of protein cquilibrium ensembles with generative deep learning. bioRxiv, pp. 2024-12, 2024.
Onsager, L. and Machlup. S. Fluctuations and irreversible processes. Phys. Rev., 91:1505-1512 Sep 1953. doi:
10.1103/hhyseev.11.55555
Liang, S., Singh, A. N., Zhu, Y., Limmer, D. T., and Yang. C. Probing reaction channels via reinforcement learning. Machine Learning: Science and Technology, 4(4):045003, 2023.
Polyak, A., Zohar, A., Brown, A., Tjandra, A., Sinha, A., Lee, A.. Vyas. A., Shi, B., Ma, C.-Y.. Chuang, C.-Y., et al. Movic gen: cast of media foundation models. arXiv preprint arXiv:2410.13720 2024.
Lindorff-Larsen, K., Piana, S., Dror, R. and Shaw, D. E. How fast-folding proteins fold. Science, 334(6055): 517-520, 2011.
Prinz, J.-.., Wu, H., Sarich, M., Keller, B., Senne, M., Held, M., Chodera, J. D., Schüttte C., and Noé, F. Markov models of molecular kinetics: Generation and validation. The Journal of chemical physics, 134(17), 2011.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint ar arXiv:2210.027444 2022.
Psenka, M., Pai, D., Raman, V., Sastry, S., and Ma, Y. Representation learning via manifold flattening and reconstruction. Journal of Machine Learning Research, 25(132):1-47,200)
Lipman, Y., Havasi, M., Holderrieth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T., Lopez-Paz, D., Ben-Hamu, H., and Gat, I. Flow matching guide and code. arXiv preprint arXiv:2412.06266 2024.
Reidenbach, D. and Krishnapriyan, A. S. Coarsenconf: Equivariant coarsening with aggregated attention for molecular conformer generation. Journal of Chemical Information and Modeling, 65(1):22-30 2024.
Machlup, S. and Onsager, L.. Fluctuations and irreversible process. ii systems with kinetic energy. Phys. Rev., 91: 1512-1515, Sep 1953. doi: 10.1103/Phy.Rev.911.5122
Rogal, J., Schneider, E., and Tuckerman, M. E. Neural- nawork-bbssed path collective variables for enhanced sampling of phase transformations. Physical Review Letters, 123(24):245701, 2019.
Maier, J. A., Martinez, C., Kasavajhala, K., Wickstrom, L., Hauser. K. E., and Simmerling, C. ff14sb: Improv- ing the accuracy of protein side chain and backbone
12

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Rojo, A. G., Rojo, A., and Bloch, A. The principle of least action: History and physics. Cambridge University Press, 2018.
collective variables for enhanced sampling of rare events, 2022.
Swendsen, R. H. and Wang. J.-S. Replica monte carlo simulation of spin-glasses. Physical review letters, 57 (21):2607, 1986.
Rose, D. C., Mair, J. F., and Garrahan, J. P. A reinforcement learning approach rare trajectory sampling. New Journal of Physics, 23(1)0113013, 2021.
Tiwary, P. and Parrinello, M. From metadynamics to dynamics. Physical review letters, 111(23):230022, 2013,
Rotskoff, G. M., Mitchell, A. R., and Eajden-Eijnden, E. Active importance sampling for variational objectives dominated by rare events: Consequences for optimization and generalization. In Mathematical and Scientific Machine Learning, pp. 757-780. PMLR, 2022.
Torrie, G. M. and Valleau, J. P. Nonphysical sampling distributions in monte carlo free-energy estimation: Umbrella sampling. Journal of computational physics, 23(2):187-199, 1977.
Schnuuing, A., Harris, C., Du, Y., Didi, K., Jamasb, A., Igashov, I., Du, W., Gomes, C., Blundell, T. L., Lio, P., et al. Structure- drug design with equivariant diffusion models. Nature Computational Science, 4(12): 899-909, 2024.
Tuckerman, M. E. Statisticul mechanics: theory and molecular simulation. Oxford university press, 2023.
Valsson, O. and Parrinello, M. Variational approach to enhanced sampling and free energy calculations. Physical review letters, 13(())090001, 2014.
Schreiner, M., Winther, O., and Olsson, S. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. Advances in Neural Information Processing Systems, 11356449-364666 2023.
Vanden-Fijnden, E. and Heymann, M. The geometric minimum action method for computing minimum energy paths. The Journal of chemical physics, 128(6), 2008.
Seong, K., Park, S., Kim, S., Kim, W. Y., and Ahn, S. Transition path sampling with improved off-policy training of diffusion path samplers, 2024.
Vanden- Eijnden, E. et al. Towards a theory of transition paths. Journal of statistical physics, 123(3):503-523, 2006.
Vanden-Eijnden, E. et al. Transition-pati theory and path-finding algorithms for the study of rare events. Annual review of physical chemistry, 61:391 2010.
Singh, A. N. and Limmer, D. T. Variational deep learning of equilibrium transition path ensembles, 2023.
Vander Meersche, Y., Cretin, G., Gheeraert, A., Gelly. J.-C., and Galochkina, T. Atlas: protein ficxibility description from atomistic molecular dynamics simulations. Nucleic acids research, 52(D1))D384 D392, 2024.
Sipka, M., Dietschreit, J. C., Grajciar, L.,, and Gómez- Bombarelli, R. Differentiable simulations for enhanced sampling of rare events. In International Conference on Machine Learning, pp. 31990-32007. PMLR, 2023.
Vincent, A connection between score matching and denoising autoencoders. Neural computation, 23(7): 1661-1674, 2011.
Sobieraj. M. and Setny, P. Granger causality analysis of chignolin folding. Journal of Chemical Theory and Computation, 1833):193-1944,, 2022. doi: 10. 10.102//acssjcte..1o0994
Xie, T., Fu, X., Ganea, O.-E., Barzilay, R., and Jaakkola, T. Crystal diffusion variational autoencoder for periodic material generation, 2022. URL https:
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32,2019. 2019.
/
Yan, Touchette, H., and Rotskoff, G. M. Learning nonequi- librium control forces to characterize dynamical phase transitions. Physical Review E, 105(():024115, 2022.
Struski, L., Sadowski, M.. Danel, T., Tabor, J., and Podolak, I. T. Feature- based interpolation and geodesics in the latent spaces of generative models. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Zeni, C., Pinsler, R., Züigner, D., Fowler, A., Horton, M.,
Fu, X., Shysheya, S., Crabbé, J., Sun, I.., Smith, J., et al. Mattergen: a generative model for inorganic materials design. arXiv preprint arXiv:2312.03663 2023.
Sultan, M. M. and Pande, V. S. Automated design of collective variables using supervised machine learning. The Journal of Chemical Physies, 149(9)944006, 2018. doi:
Zheng, S., He, J., Liu, C., Shi, Y., Lu, Z., Feng, W., Ju, F., Wang, J., Zhu, J., Min, Y., et al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, pp. 1-10, 2024.
Sun, L., Vandermause, J., Batzner, S., Xic, Y., Clark, D.,
Chen, W., and Kozinsky, B. Multitask machine learning of
13

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlug Functional
Sípka, M., Erlebach, A., and Grajciar, L.. Constructing collec-
tive variables using invariant learned learnede representations. Jour- nal of Chemical Theory and Computation, 1933):877901, 2023. ISSN 1549-9618. doi: 10.100///ccs
14

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
A. Extended Related Work: Transition Path Sampling
rich landscape of tools has been developed for TPS, for which we refer to existing surveys for a more exhaustive description (Bolhuis al., 2002: Dellago et al., 2002; Vanden- Majden- al., 2010). Traditional shooting methods perturb the initial or intermediate states of a known trajectory to generate new trajectories via Metropolis-Hastin Hasting criterion (Mullen et al., 2015; Borrero & Dellago, 2016; Jung et al., 2017; Bolhuis & Swenson, 2021). These often suffer from high rejection rates, correlated samples, and the need for expensive molecular dynamics (MD) simulations during sampling. Another class of methods is based on adding an adjustable biasing potential to enhance the sampling of slow events, which includes umbrella sampling (Torrie& & Valleau, 1977), metadynamics (Laio & Parrinello, 2002a), and more advanced techniques such as eABF (Darve & Pohorille, 2001). These approaches require a carefully constructed, low-dimensional mapping of the problem via collective variables (CVs), which can be challenging, particularly when the characterization of the system around the transition state uncertain, Attempts were made design CVs with ML methods (Sultan & Pande, 2018; Rogal et al., 2019; Chen & Ferguson, 2018; Sun et al., 2022; Sípka et al., 2023), yet they remain challenge for many-atom systems. ML approaches have also been used to learn the biasing potential directly, including approaches based on stochastic optimal control (Holdijk al., 2024; Yan et al., 2022), differentiable simulation (Sipka et al., 2023), reinforcement learning (Das et al., 2021; Rose et al., 2021; Singh & Limmer, 2023; Seong et al., 2024; Liang et al., 2023), and l-transform learning (Singh & Limmer, 2023; Du et al., 2024). In all of these approaches, unlimited access to the underlying potential energy and force field are assumed, but samples from the underlying data distribution are not available. As a result, the methods must be retrained from scratch for every new system of interest. Additionally, expensive, simulation-bart based training procedures are often employed, limited scalability to larger systems. Interpolation-basee methods, such the Nudged Elastic Band (NEB) method (Henkelman al., 2000) and the spring method (Dellago al., 1998), introduce springs between images construct transition paths (spring method) or to directly locate saddle points (NEB). However, significant challenge for both approaches lies in generating an initial guess, which is inherently unknown a priori. Among interpolation-like approaches, the Onsager-Machlup (OM) action has also been explored for TPS. Due to the limited availability of automatic differentiation techniques at the time, Laplacian operators were consistently avoided. This restriction limited its application to very low-dimensional problems (Vanden-Eijndee & Heymann, 2008; Fujisaki al., 2010), or led to the development of Laplace-free action formulations (Lee et al., 2017).
B. Proofs for Related Generative Model Objectives
For the sake of readability, we replicate proofs showing that both the training objectives for DDPM and flow matching models are equivalent to training against the score of a noised version of the data distribution (or in the case of flow matching, an invertible polynomial transformation of this score). See Vincent (201) 1) and Lipman et al. (2024) for example proofs for DDPM and flow matching respectively.
B.1. A note on the DDPM reverse process
The sampling process for a DDPM can be written as the terminal condition (0)) of the following process:
x (O..T)
(22)
x(
1
(1) +
8r so (x i)) + v B,z.
(23)
V1-3,
V1 V1-á,
While this process is definable as an Guler-Maruyann discretization of an SDE, is not well suited for optimization over trajectories over the data distribution, since the vector field su (xx(), is changing throughout the trajectory, and iterates near the noise T= will not necessarily follow dynamics determined by the data distribution.
B.2. DDPM and score matching
Below a proof for the equivalence of standard DDPM training to score matching. Theorem B.1 (DDPM-Score Matching Equivalence). Let Paat ( be the data distribution, and let x., be the noised variable defined through the forward process:
Ir = V0r2o+V1--0 T~ T=Uni((1,,...Ta))
Io ~ Ddaus E~N(),T);
(24)
where ā, (0.1). Letp, Le, )= = ( (2a: (1-a ) )dzo be the marginal distribution x, Then the DDPM objective, defined as the following:
LDDNM
del B [[k- [1c--co(=,:7) (x(
(25)
15

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
satisfies the following equality:
CDDPM( = =r.= [1 (1=, (2+) so s0 +C,
(26)
where so (T( del co (z( vV...r/////aaaaa and C is constant independent of
Proof. We begin with the DDPM training objective:
CDDPM = =Er (0)
(27)
The core of the desired result is Tweedie's formula, which relates Gaussian-based denoising to the score of the noised distribution. For any random variable 2 generated as + an where is an arbitrary random vector, ~ (N())) and a > O, Tweedie's s formula gives the following posterior expectation:
Eiu|2] = = =o====+/VViie (olgg()(())
(28)
where p(z) = of (p(p)d is the full marginal distribution of z, the forward process, o, generated via 3, - va=var= - which corresponds to:
μ== Vā
== =
(29)
Here, is a random variable (dependent on 2o), not fixed parameter. Applying Tweedie's formula to the marginal distribution Pr (x+),, we obtain:
E Vā 2. = (1- logp, (x+). Dividing through by var gets the posterior of the original sample zn:
(30)
Tr 1-a, Va, (o7))(x7). Va- Var
[[xa|==>==
(31)
From the forward process definition, rewrite in terms of
(32)
ā
and take conditional expectations given x to get the following:
Ble| - Fr
v V E
F x-
(33)
V1-ā,
Er Var
+ L-at V., logp (IT)
Va.
va,
(34)
a+
2G JT. (1 (1-ā;) V., logp,(r7)
(35)
Vl-ā,
1 aa V., lo.po (xy).
(36)
Using the law of total expectation, we can expand the DDPM loss conditioned on Wr
(8) [L [ (0
(37)
For any random vector g. [B[lE c|l2] is minimized when c = E[g]. We can then make the following bias variance decomposition:
E. [|le - co (Kr,7 ) | = ||[|e| [ (x )1 ( w - [T,7]. Since the variance term is independent of substituting EE[c x and factoring out / a, leads to:
(38)
(D) (0) (1 (1-ār) V logp (-
()( (-co(on.7)
+C.
(39)
VI 11
By defining s0(x+,):=-<o(f - (0(x),7)/ - we obtain the score matching objective.
16

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
B.3. Flow matching and score matching
We now provide proof for the flow matching setting, showing that the training objective is also similar score matching objective, with simple transformation between the flow matching targets and the scores.
Theorem B.2 (Flow Matching Score Matching Conversion). Let Páaa(Fo) be the data distribution, and let be the noised variable defined through the interpolation process:
2+ - -0- 10,11, T~Unif((O,1), 21 Pdata Fo (0,T)
(40)
where ar ( -0.1000000010100 [0,1]
are strictly increasing and decreasing functions respectively that satisfy ogg =0 =0= a1 00 = 1. [ Paura (2,;0011100 Idroo be the marginal distribution of E.. Then the flow matching objective, defined
Let pr
(z,)= (r(()
as the following:
(FM( det
Er 20. / (x), ( )12)
(41)
where u. (1( x + t0 the instance-wise curve velocity. The flow matching objective then satisfies the following equalities:
1. We
equivalently
train against
targets of the unconditional velocities
Wr ur(x)
Exo Exo (xx1 x= 1ax1a110t0000 +=(x0:
(FM(0)= =
(x+)) (1) +C]
(42)
where C' is some constant independent of and
2. The equality Vzlogp- = dz d, a. u (x) holds, allowing us to write the flow matching objective in terms
of the score:
(FM(0)= = wo Er,
[
år o ,, V.. logp (IT)
12 +C.
(43)
OtT
0rr
Proof. For part 1, we can expand the flow matching loss integrand by telescoping with respect to (I:7)
u, 2 (aaall//lll_a::::000
+u, (xo.
(44)
= ||u - ( u, (x. u, (x,) - (1 (2r
)11
(45)
Since =, ||u- (x+) -u, or (2 )l 2 is constant with respect to 0, it suffices to show the following:
u, Ey E7,20.2, (4) (x+) u, (x), =0.
(46)
Note that by definition we have the following relation between u and
E
[e, (20, x [, [ Φ1 +0, x, [,, - u (I+).
(47)
20.21
20.00
We can then write the following by expanding the expectation using the tower rule:
Er,za. (Mo(I,,5) - (.xr).,1 (x,, (To = EE
(x,
E
[ (uo(xTT,T) (xx,), ), , (xx, ) v. (zo [xr.7]
(48)
E [ua((TTT) (x (W))),) (x+) -Erg [Vr [
(49)
= E (( me.
(50)
(51)
This concludes part Note that the interpolations x, = 1 20 also follow proper form for Tweedie's formula, allowing us write the following:
= xr + lloog (I,r), E Ela,,110,
(52)
- Er + o²
lo (I7)
(53)
Or
17

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlug Functional
Noting that To =0-0,
we write u, as the following:
dr
ur (r(x))
E
[áta +
xTT,
(54)
#0:21
di I
[2]
-
+
E
[r0 2i - xx,
(55)
20,21
20.21
o x+
E
[1 24 = IT
(56)
Tr
20,21
x+ o
V (o))pp((x))
xt
(57)
d
0,0+07
r
V (zlogp((r).
(58)
Qr
This proves the desired relation between u, and Vlogp,, and plugging into part achieves the desired flow matching loss equality.
C. Derivation of the Onsager-Machlup action
C.1. Overdamped Langevin dynamics
We start the description of our system by formulating the well-known Hamilton equations. The variables we are solving are (t): IR+ + IR and the corresponding momenta (t))) IR + Rd with a constant vector m, representing the mass of every particle i in the system. Hamiltonian equations are formulated as follows
xt) Pt(t)
mh
(59)
px(1) aU (x((())
Ox;
While these equations maintain energy and contain representation of temperature, modified SDE, with the term W(t) representing a Wiener process and damping constant y
(t(t) xt())
m
(60)
pa(1) aU(((())
(1(1)
+ + TW())
Ox;
mi
or equivalently in one second order equation:
au (x(()))
m, x. (1)
7 + 2m))kg7W((),
(61)
0x:
represent system that experiences thermal fluctuation. Although the original Hamiltonian system is trapped in an energy well forever, the one guided by Langevin dynamics may overcome barriers between wells in finite time.
A question then arises. Of all the possible paths of fixed physical parameters and time that connect two minima, which is the most probable? How do we calculate probabilities and penalize high energy regions or paths that are making too large steps? The answer is provided by Onsager and Machlup in their works (Onsager & Machlup, 1953; Machlup & Onsager, 1953). The second reference handles the full equation Eq. (60), while the first is a reduction to so-called overdamped state where R&T
the term x(t) can be neglected. After introduction of two auxiliary vector quantities G = m:y and D;
we get the form Ç:
1 (x)(x(
x;
+ 2DD,
(62)
Si
0x:
or equally just with F(x(()) (x(())))
8x,
1 F(xx()++
(63)
So
18

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Functional
Further, we will follow a more general setting of the Langevin equation consistent with Eq. (1). To recall:
x vx))
(64)
and
((x): = Uxx) 4 (x)
(65)
(66)
C.2. Most probable path under Langevin dynamies
Considering a single particle (for more particle systems see e.g. (Kappler Adhikari, 2020)), since Eq. (I) a stochastic differential equation,
also write partial differential equation for the probability density of the particle guided by these equations. In this case it is well-known Fokker- Planck equation (note 2 of vector will be understood as divergence operator to save space)
8x
OP
D
(67)
Ox
Ox
Ux
As we will consider only potential forces in this work, let us denote two parts.
(xx) = 86(x)
Now we will split the derivation into
1. + 0:
Let us consider the solution in the following form:
P(x,t xo) (00== -0x-x002
4D4
(68)
Then for the sequence of points in space and time (x' .t') ,t'),(x²+
(xx"))tt))))) (x..). (x) we can write the following probability:
N
P(x
Xx 1
P( (xo).
(69)
j=1
Let us denote // 11-1i-1= as we pass through continuum limit in time. The probability can be rewritten by plugging in
solution Eq. (68) into
N
D4) (
xo)
BN exp
Σ
[xj
x 1)2
(70)
4Dc
1=1
j=1
To make sure we can pass into the limit let us rewrite
Ine
(71)
We now focus the argument of the exp function. We can modify it to the form
(x--)--)c
X x
2 c.
(72)
4D
e
By passing into the limit N - oc and realizing that epsilon can be rewritten by its definition to c= N. we get the following integral form:
1 dt.
(73)
4D
Jo
Note however, using the identity a blaa
blua the first part of the product goes to infinity:
-N
t
lim
4D
(74)
N
N
19

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machilee Functional
evaluation of this limit directly would be too hasty. One must consider the probability derived in the broader context of integration across the path. In that case, the constant will serve to normalize the probability. The fact that it does not depend on also means that the probability of the path does not, relative other paths, depend on this prefactor, and only the exponential part is important. To find out more about precise mathematical justifications, we refer the reader to (Gel" fand & Yaglom, 1960). We shall denote the constant before exponential as C from now on because, as is not dependent on x, it will not influence our calculations. The final probability of a path is then given as follows:
(-
1 J J (x (x))))
P(x.t)=Cexpp
(75)
To maximize the likelihood of the path we clearly need to minimize the action
1 1 (x(s))
So(x(t)) - 4D
(76)
Intuitively, the most probable path under no drift is the one that does not move from it's origin. The longer the trajectory, the less probable is.
2. 4 O0:
We will recall the assumption + Vo((x) and shall use transformation
1
P(x,t] (6) - )exp
(ss)ds
(77)
2DC x(0)
where from the properties of a potential function we can evaluate the integral to
(xx) def 1
x(t) 4 (s)ds=
x(1) (ss)
(- (-x(()+((x)))
(78)
2DC
2DC x(0)
2D
x(0)
So for clarity:
G(x,t) -xxx)
(79)
( P(x,))e,
(80)
1
Vo(x)= = 2D VΦ(x).
(81)
We then plug this transformed function into Eq. (67). We will now derive the equation that G(x,() has to fulfill. Let us evaluate left-hand side of the Eq. (67)
0P(xt) aG(x,t) -(x)
(82)
For the right-hand side lets evaluate first the term:
a 86(x) P(xxt)
gx)(
1 10P(x,t) 0o(x)
P(x,t)
dx
< dx2
5
dx
dx
= =>2DPxx,t)
8
2DD
Px(x,)
(83)
0x2
0x
0x'
-DDG x(x) 2
2D
0G
x
+2DP
O
2
0x
0x
dx
0x
While the other term can be written as follows:
0p(x,t)
0G(x,t)
- -G((xx)
-70
Ox
Ox
Ox
(84)
OG(x,t)
- -P(xtt
00
dx
0x
20

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
The second derivative then with function arguments omitted for brevity, yet remaining the same
D P
=D
G - Ox?
-D
0G
-
-D
aP
06
DP
86
Ox2
8x
dx
Üx
ox2'
(85)
G²G -
0G
05
dΦ
2
8
=D 0x?
C
2D
e
DP
DGe
+
0x
δx
0x
0x?
After subtracting the terms the right-hand side, we get the following:
aG ক -D J
- DGe -
0Φ
2
+De D
-50
(86)
at 0x
Ox
8x2
Or written nicely after exponential cancels and we return ((x) from /
(G(x,t)
(x(()x)) G(x.t)
[6(x)
Φxx
D
(87)
Øt
0x2
4D
<
0x
2
0x2
This is a well studied diffusion-reaction equation
0u(x,1) D
Pu((x,)) ku((x,))-
(88)
Øt
Ox2
Notice for pxx) =0 we already solved this equation as is identical Fokker-Plance where F=0. Let us call this solution ug. Another observation is that for this equation we can formulate solution in the form
2xx,t) = uo (x, (okdd
(89)
2012
- Ce 4D1
where C is some arbitrary normalization constant as in the previous solution:
-(x-xo) x(1)
G(x,t)=Coxpp
- k(s)ds
(90)
4De
x(0)
and the original f(x.t) using Eq. (77):
xxx
px(1)
(-6(x(6))+
1 2DC (ss) d
P(x,t)=Cexp
(91)
4Dt
Now we repeat the same multiplication of probabilities for small time increments. However, this time, the situation is easier as integrals would simply extend the sum. Therefore the only limit would be in the first term exactly as done before. The final probability of the path is then as follows:
4D 4D / (x)2+
1 0á
2D
2000
P(x,t)=Cexp
24
(92)
( 0x
0x2
The negative argument the exponential will again be an action to minimize:
S(x(t))
4D
1
(x)2+
( 0b
( (xx
2
2D
2D
0xx
-
(93)
This can be further modified, by integrating forces along the path and using forces instead of potential, to the more common form:
1 - (d(x) 4D. x²
1 00
2D 8 ds.
S(x(1))
(94)
2D
Ox2
This procedure derive the Onsager-Machlup action is similar that in (Mauri, 2012).
21

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
D. Fixed endpoints
For the entirety of this paper, we operate with fixed endpoints. This means the actual minimized action will be reduced simply to the following:
106 2
2D 8
((<(())=<T)) x²
x0x Ç 0x2 ds.
(95)
The first and simplest strategy keep containt constant is include a penalty in the form
= (x( - ) 2 -x))
(x(t)-xx)))
(96)
Interestingly, as verified experimentally, the penalty term effectively works the same as using more simple and straightforward method. The method of choice was to set the endpoint gradients to 0 manually, Only couple of points will be affected, and the majority of the trajectory is the same for both approaches.
E. Numerical discretization
Another aspect consider is the numerical discretization of the action. (Adib, 2008) discusses the different numerical evaluations of the action stemming from the stochastic nature of the Langevin equation. Namely, the Ensager-Machlyg action depends on the discretization convention used for the SDE (Itô or Stratonovich). We consider the Stratonovich convention for this paper, which yields the following Onsager-Machlog action:
1
N
15-xx-111 x
((x,))
2DAt +
S(Xo,XI..Xn)) =
+ +tt
V.
(97)
4D
At
j=1
If the Itô convention was used, there would not be a Jacobian term v (Cugliandolo & Lecomte, 2017).
To make the multidimensional, multiparticle system discretization clear, we extend the sum along spatial dimensions (index j) and we also sum particles (index k). The coefficient Ç is now a vector since it is originally CA = where My is the vector masses. The total action is then,
I-11 1 No ΣΣ ADA xx,
200
At 4DC3
At v. xx(()x 1 10
SIx (0)
(1)
(i+1)
(98)
26,
F. Classical Force Fields on All-Atom Proteins
We demonstrate that our OM action optimization framework is broadly useful for transition path sampling even beyond the setting of generative modeling. Specifically, we aim to find all-atom transition paths between the unfolded and folded states of the protein Chignolin and Trp-Cage, using a differentiable Py Toroh implementation (Doerr et al., 2020; Sipka al., 2023) of the Amber.ffl 4f et al., 2015) forcefield and the TIP3P implicit water model. We choose the physical parameters of the OM action to be consistent with commonly used values in molecular simulations (see Table 2). Since we do not have generative model from which obtain initial path guess via latent interpolation as described in 4, we instead employ hierarchical unwrapping warm-up procedure described in the Appendix G to obtain initial paths. As the classical force field is dominated by quadratic terms whose I Laplacian is constant and thus uninformative for optimization, we use a zero-temperature approximation and optimize with the Truncated OM action. Using the Truncated action, we obtain a physical transition path of length 2.6ps (shown in Fig. 6). This much lower than previously reported transition path lengths for Chignolin (Sobieraj & Setny, 2022: Lindorff- Landorff-Laren et al., 201 1). which can be explained by the fact that our trajectories proceed between the target states without fluctuations that would occur in unbiased simulations. The entire optimization took the order of hours one NVIDIA RTX A6000 GPU, including the generation of initial trajectory.
G. Additional Details on Onsager-Machlup Action Minimization Method
Initial Path Guess Methods. We provide a complete description and algorithmic formulation of the initial path guess method using a generative model, mentioned in Section 4.
22

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
Chignolin
Trp Cage
Time (ps)
0.0
0.5
1.0
1.5
2,0
25
Figure 6. Transition paths from OM optimization of all-atom chignolin and trp-cage with classical force field.
Table 2. Hyperparameters used for OM action optimization all atom proteins
Hyperparameder
Chignalin warmup
Chignolin
TRP Cage warmup
TRP Tage
Number of points per path
40 2600
2600
40 2600
2600
Action Type
Truncated
Truncated
Truncated Adam
Truncated
Optimizer
Learning Rate
10-4
10-3
10-4
10-b
Action Timestep (At)
fs
fs
1ff
lifs
Action Friction (*)
10 ps
10 ps
ps
10 ps
Formally, consider generative model with a non- non-parametric, probabilistic encoding (i.e corruption) process g(x, (xo). and a corresponding, parametric decoding (i.e generative) process po (xu x,)). We first roto-translationally align the endpoints of the path via a Kabsch alignment (Kabsch, 1993). We encode the aligned endpoints of the path into the chosen latent level for the
initial guess, Tinitial to produce two latent endpoints (0)) (Z (0) and (L)) qz (Z (z))). ). We then interpolate linearly
generate latent path Z (z(()) (1- )) + i (z) =Ejo,,] We can then either decode the path back to the confi configurational
space via po (x) to obtain path x, in which the subsequent OM optimization would occur in configurational space, or defer decoding, in which case optimization occurs at the latent level Tinitill starting from the latent path Z. Intuitively, larger values of Tintiall produce more diverse initial guesses at the expense of decreasing correspondence to the endpoint states.
Algorithm 2 Initial Guess Path Generation with a Generative Model
1: Given generative model with non-parametric encoder y. and decoder po-
2: Function InitialGuess (x (0) (L.) .Tiitial)
3: Align both samples via Kabsch alignment:
4: xU)) (_ KabschAAignll (xign(((()) (0) (x)
5: Encode both samples into latent level Tinntal of the generative model:
6: (0) (Z (9), (4) ~ (Z (x))
7: Interpolate linearly (or spherically) in the latent space to generate an initial guess latent path:
8: Z= Z= (= = (1 - (z) +
+ 22(4) ]ee(o.z]
9: Decode each point the initial latent path Z from Timiial to = to produce data path:
10:
x= [x(+) ~ (x(22(( )
Return
When using classical FF, we do not have access to generative model. Thus, we must use different scheme than what is described in 4 to compute the initial guess path. We first start with small number of replicas in each basin, creating a large gap in the middle of the path. We then optimize with an unphysically large path term, creating a short but interpolating trajectory. After we are satisfied with the initial guess we multiply each replica twice, creating a path of twice the length that we then optimize again. This simple procedure is repeated until we reach desired length of the path. This procedure is described in Algorithm 3,
Hutchinson Trace Estimator. The third term in Eq. (14) the trace of the Jacobian of the force, F Fopp): or equivalently for conservative forces, the trace of the Hessian (Laplacian) of a scalar energy. Naively computing
23

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Machlup Functional
Algorithm 3 Initial Guess Path Generation with Iterative Unwrapping
1: Function Inesial xx!" _I.NN)
2: Initialize trajectory by copying boundary points L1 times on both ends.
3: x x= =xx(0), ,....x(() xx2),,,.. (x.,,,,,
4: For m from N repeat:
5: Duplicate every point along the path: x {x(0) xx
x(1) (1),,,(())((
6: Optimize Truncated Onsager- -aaphlu action: argmin mn (x)
7: Obtain initial guess x=
x x()) ,...,!!?" .h (1)
gradients of this quantity can be prohibitively expensive. We thus employ the Hutchinson trace estimator (Hutchinson, 1989) to accelerate computation. Formally, let H(xx V (x, E RR N We approximate the trace of as tr(H(x)))= N -E=IVTH(()) (x)())vv where v (0,T). By leveraging vector-Jacobian products (VJP), we can compute the trace without materializing H or its diagonal elements.
In practical terms, the estimator converges rather slowly. Let us denote the approximated trace by Tr. One can derive the variance of the estimator as,
Var(Tr)= 1 Var(v, ((x))),
(99)
N
which, when v, are distributed identically means the error of the trace estimator decays
C
I (T=TT(H(x))s
(100)
VN
Where C depends on the properties of the matrix. This convergence rather slow and means that one requires many iterations to arrive to accurate value of the trace. In practice, however, found that N = 15 worked well and led to smooth OM optimization. This is likely due to the fact that our trajectories composed of many neighboring points that likely had similar Laplacian values.
Selection of optimization time. As described Section 4, the time Tops used to condition the generative model score function sp (o, is treated as a pyperparameter. In principle, Topx = ensures maximal correspondence with the true atomistic force field for Boltzmann-distrbuuted data, but consistent with Arts et al. (2023), we find in practice that a small, nonzero value works better. In Fig. 7, we show the average cosine similarity between the true forces and our pretrained denoising diffusion model's s score function various values of Topi for the Müller-Brown and alanine dipeptide systems. Notably, Topt = 0 is not the optimal time conditioning with respect to true force recovery for DDPM.
H. Müller-BBBrB Browen Potential Experiments
We provide further details on the Müller-Brown experiments Section 5.1. All experiments were performed on a single NVIDIA RTX A6000 GPU.
Potential Parameters.
The exact form of the potential used is the following:
U(x.)) 17.3e -.000000--4433-000009
-8.7e - 3.72-00000000-3333--000001-1
- 11
+ 1..c 0.0
This generates the potential shown in Fig. 1. Fig. 8 shows OM optimization using the analytical potential as the force field. Increasing the diffusivity yields paths that cross higher energy barriers, aligning with physical intuition. The results with the diffusion model in Section 5.1 align with the paths derived from the analytical potential, confirming the validity of the diffusion model as an approximation of the forces.
24

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
0.950
Cosine Similarity
Cosine Similarity
Predicted Force vs. .Tue Force
Predicted Force True Force
0.925
0.900
0.88
0.875
0.84
16
Diffusion Timestep
Diffusion Timestep
(a) Müller- Brown.
(b) Alanine Dipeptide.
Figure Cosine similarity between true force fields and learned force field from DDPM's learned time-dependent model so (--Tog) over different diffusion latent times Tapt- Results are averaged over all paths and atoms, for a Müiler-Browa potential and Alanine Dipeptide. Notably, Tap is not the optimal time conditioning with respect force recovery for DDPM.
50
OM Optimization Diffusivity-0000
40
OM Optimization Diffusivity
OM Optimization Diffusivity=4.0
30
20
10
9
20
30
40
50
Figure 8. OM optimization can be done with analytical potential. We show paths generated with OM optimization using the analytical potential with a timestep of 1, a friction of 1, and multiple diffusivities. Higher diffusivities, corresponding to higher temperatures in physical interpretation, cross higher energy barriers, aligning with physical intuition.
Data generation. To ensure that the transition region adequately represented with relatively short simulation times, we choose initial conditions for the simulations by uniformly sampling the transition path resulting from OM optimization under the true MB potential. We generate training data by running unbiased, constant- -esstaatteempeatuuue simulations with the MB potential under Langevin dynamics. We 1,000 parallel simulations for 1,000 steps, yielding total of million datapoints. Of these, 800,000 are used for training, and 200,000 are reserved for validation.
Training. We then train a standard denoising diffusion model on this dataset, with the denoising model parameterized by a 3-yyer MLP with a GELU activation (Hendrycks Gimpel, 2023) and hidden dimension of 256. The model is trained for epochs, with a batch size of 4096 and learning rate of 1e using the Adam optimizer (Kingma & Ba, 2017).
25

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machiee Functional
50
L.D. Samples
OM Interpolation
40
30
20
10
90
20
30
40
50
Figure 9. OM optimization from flow matching model. The analogous experiment Fig. 8, but using flow matching model trained on Müller-Brown data, and using the extracted score via Eq. (17) for OM optimization. Since the stochastic encoding/decoding processes of flow matching models deterministic and the Müllee-Brown setting has single minimal landscape, the encoding/decoding scheme itself cannot generate diversity in this setting. The trajectory generated at diffusivity D -=10
Energy Laplacian Term. We estimate the Laplacian of the potential energy surface by using the Hutchinson Trace Estimator (see Section G). As shown in Fig. 10, one random vector (N = 1) is enough to capture the local minima and the energy barrier using the Hutchinson Trace Estimator, so we use N w=1 in our experiments. Using more random vectors gives a less noisy estimate the Laplacian, trading off accuracy for computational expense.
Analytical Laplace
True Diffusion Laplace
Hutchinson Diffusion Laplace n=1
Hutchinson Diffusion Laplace n=5
40
40
40
40
>, 20
>
20
20
20
%
%
%
%
25
50
25
50
25
50
25
50
Figure 10. Hutchinson Trace Estimator accurately estimates the Laplacian. The difffusion model learns an estimate of the Laplacian that captures düller-Browe energy wells. The Hutchinson Trace Estimator efficiently approximates this Laplacian, and estimate hecomes less noisy when using more random samples.
OM Optimization Details. We pick two points on the potential energy surface (PES) at alternate ends of the transition barrier as target points for interpolation. The hyperparameters for the OM optimization are given in Table 3
Committor Function and Rate Estimation We provide additional details on the committor function and transition rate estimation experiment presented in Section 5.1.
According to transition path theory (Vanden-Fijnden et al., 2006), for transition event between endpoints A,B E Q the committor function q(x), captures the probability that trajectory initiated at xo = reaches before AA:
q(x) =()x<() |xo xx];
argmin x (00+)
{x+ AUB::xx0000
(101)
where hg is the indicator function for reaching state B. The transition state ensemble is formally defined as the level set
26

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Table 3. Hyperparameters used for OM action optimization the Müller- Brown potential with diffusion.
Hyperruraametrr
Number of Generated Paths
50
Action Type
Full
Tnitial Guess Time (Taiil)
8
Oplimization Time (Tip.)
Optimization Steps
200
Optimizer
Adam
Learning Rate
0.2
Path Length (Z)
50
Action Timestep (A.))
0.01
Action Priction
0.01
Action Diffusivity (D)
0.0.0.4.0
(xx)):(x)=00.55 0.5). The committor function (xx) can be used to estimate transition rates between a reactant state A and product state B using the following formula (Vanden-Eijnder et al., 2006):
V do KnT T ( (>)((x))
k:
(102)
XEI
Y
where p(x) the probability density under the Boltzmann distribution and Q the configuration space. Thus, with a differentiable estimate of the committor function, we can compute reaction rates as an ensemble average over samples from p(x). By applying the Backward Kolmogorov Equation (BKE) and Vainberg's Theorem, we can show that the committor function is the solution to the following functional optimization problem (Hasyim et al., 2022):
((x) = argmin = 1
V Vx](x)
(103)
q 2
9
subject to the boundary conditions q(x) = ( = the committor function qa (x) by extremizing the following loss function:
ØB. Thus, we can train a neural network approximation to
1 ((0)= 2 V (x) (x) +2 (90(x))))++)) AB 1 (11-))(x)))))
1
(104)
The first term minimizes the BKE functional, and the second and third terms enforce the boundary conditions with penalty strengths AA and Äg. respectively. The ensemble averages in the loss, which are high-dimensional integrals whose computational cost grows exponentially with system size, are in practice by replaced by aepportace-samplle Monte Carlo averages over a dataset of samples Drrain = [ from MD simulations, MCMC, or any enhanced sampling method (such as OM optimization). Once trained, the neural network committor function go can be used to estimate the rate via Eqn 102, substituting qe (x) in place of gxx), and averaging over Durin instead of the intractable full ensemble average over Q
KBT (IVxgo Exx V (x)
kg
(105)
Notice that the neural network estimate the rate kg is eqquivallnt (up to constants) to the BKE functional loss C[j]. The primary challenge with learning the committor function in this way is the issue of sampling: the optimization problem is dominated by rare events with large values of Vx4o (x){2 2 (i.e events in the transition region). Therefore, if the transition region(s) are insufficiently sampled, the training procedure will likely fail to converge, and the estimate rate ky will be inaccurate. Therefore, committor and rate estimation using this method reduces to rare event sampling problem. Inspired by Hasyim et al. (2022), our idea is to use the transition paths obtained from our OM optimization procedure (shown in Fig. 1 Ia) as samples over which to minimize Eq. (104). Specifically, for the Müller- -Brwnn potential, we first sampled more points by initiating unbiased Langevin dynamics simulations along the OM-optimized paths (100 simulations, each 50 ps long, k T=,T=)),, 1), using the diffusion model's score function at 7 0 as the force field (we could have also used the analytical forces from the Küilerr-roow potential). This resulted in collection of points Pirain (a;) N (shown in Fig. 1lb). To account for the biased initial conditions of the simulations, reweight the samples in Dunin to the underlying Boltzmann distribution using the reweighting factors
W:
POM (122)
(106)
Σ,
PoM (oM(22)) (222)
27

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Cantciciooouuntinnan
where 3 T (x) is the analytical Müller-Brown potential energy, and Pom is the empirical density of Purain obtained via binning. We train 5-layer MLP with hidden dimension of 64 and sigmoid activation function approximate the committor function. The model is trained for 2000 steps with batch size of 4096, using Adam optimizer with a learning rate of 0.0001 and cosine decay scheduler. The boundary condition loss weights are A = λB 20. The final estimated committor function profile is shown in Fig. 1lc. We obtain a transition rate estimate of 1.3 x 10-5, compared with the true rate of 5.4 x 10- obtained by solving the Backward Kolmogorov Equation numerically using a finite-clement method in FEniCS (Baratta et al., 2023).
(a)
Initial OM Path
(b) MD Sampling From Path
Estimated
(d)
True Committor Function
(c)
Committor Function
Figure 11. Efficient committor and transition rate estimation on the 2D Müller-Brown potential using OM optimization. (a) Initial OM optimized transition path with pretrained denoising diffusion model, using Tept = 4 and 100 paths. (b) Subsequent diversification of the sampled transition paths by performing unbiased Langevin dynamics simulations initiated along the OM optimized paths (100 simulations, each 10,000 steps, TT = 1) with the difffusion model's 's score function at 7 = 0. (c) Predicted values of the committor function (x). obtained via minimizing ( V qt (x) 1),, where (.)a denotes ensemble average over the samples collected in part b) and the committor function ge (x) is parameterized by 5 layer MLP with hidden dimension of 64 and sigmoid activation. (d) True committor function calculated by solving the Backward Kolmogorov Equation numerically using a finite-element method in FEniCS. The reaction ls computed via the relation k kn (x) ))- Our predicted reaction rate is 1.3 x 10 5, compared with the true rate of 5.38 x 107 5.
I. Alanine Dipeptide Experiments.
We present additional results on the alanine dipeptide results in Section 5.2.
Details on Traditional Enhanced Sampling Baselines. For MCMC, we report metrics found in et 2024) for the fixed- length, two-way, exact method. We implement metadynamics ourselves with Ramachadran angles as CVs. We use hill height of 0.lkcal/mol,, a width of ι0Å, and a temperature of 300K. We report the number of force field (or generative model score) eval- uations required per 1,000 step transition path (a batched computation over many samples is considered to be one evaluation).
28

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Diffusion Model Training. We construct a training dataset from MD simulations of alanine dipeptide 1000K (a large temperature is chosen to promote the inclusion of transition regions in the data). We then train a diffusion model these samples, using the Graph Transformer architecture described in Section We train models for 50,000 steps, using Adam optimizer with a learning rate of 0.0004 and a cosine annealing schedule reducing to a minimum learning rate of 0.00001. We use an exponential moving average with o o==00=000 0.995. The diffusion model uses 1,000 integration steps at inference time.
OM Optimization Details. We choose a path length of = and noise/denoise by 300 steps (out of 1,000) to initialize via Algorithm 2, then conduct OM optimization for 50,000 steps, fixing latent diffusion time Topt =3. 7 and At are set to 0.7 and le-4 respectively. After OM optimization, we run short energy minimization procedure for 1400 steps using the Amber f classical force-field to relax structures, followed by OM optimization using BFGS algorithm that required around 21,000 steps.
Calculating Free Energy Barriers with Umbrella Sampling. Since we optimized with the truncated OM action Eq. (4) we obtained minimum energy paths. To compute free energy barriers that are necessary for reaction rate predictions and comparisons with literature, we employ some sampling to estimate the entropic contribution. On alanine dipeptide, we use the following procedure to obtain free energy profiles along the OM-optimized paths.
1. Perform umbrella sampling (Torrie& & Valleau, 1977): short, 10ps, simulations with a harmonic string potential centered around the points that define the original minimum energy path. The spring constant used had value of 5 kcal/ //lo//ol//// rrad
2. Fit a collective variable (CV) on the resulting samples. We choose Principal Component Analysis (PCA).
3. Employ the WHAM algorithm (Kumar al., 1992) to obtain the reweighted potential of mean force, shown in Fig.
J. Fast-Folding Protein Experiments
We present additional details on the fast-folding protein results in Section 5.3, All experiments were performed on a single NVIDIA RTX A6000 GPU.
MD Simulations with Diffusion Model To assess the computational efficiency and accuracy of our OM optimization relative to alternative methods of sampling the configurational space (see Fig. 4a and b), we rum Langevin MD simulations for varying lengths of time (up to 12 ns) using the diffusion model's score function as effective force field (Arts et al., 2023), with a timestep of 2 fs. To ensure a fair comparison, we set the number of parallel MD trajectories to be the same as the number of transition paths sampled with OM optimization.
Markov State Model Construction. We provide further details on the Markov State Model analysis used to evaluate the qual- ity of transition paths for the fast-folding protein experiments. We largely follow the procedure described in Jing et al. (2024b).
We perform k- k-means clustering of the reference MD simulations into 20 clusters using the top 2 Time Independent Component (TIC) dimensions, which are fit on the pairwise distances and dihedral angles of the protein configurations. We then fit a Markov State Model (MSM) with a lagtime of 200ps (the frequency which the simulations were saved) to obtain transition probability matrix T between the 20 discrete states in the MSM (e.g, T}, p(s klss =j). where st and $1+1 are the states at time f and ++ 1). This constitutes the reference MSM.
To evaluate transition paths sampled from our OM optimization method, we first discretize them under the reference MSM (i.e represent them as a sequence of cluster indices between and 20). We subsample the paths to be of length 20. From this, we compute the probability of the path under the reference MSM via:
We also sample 1,000 discrete, reference paths of length L L= 20 (corresponding transition time of 200ps 20 4ns) from the reference MSM, conditioned on the start and end states $1 and s, (these are the cluster indices of the transition endpoints ()) and x4))). This can be achieved by sampling states $1...899 iteratively as
T-L-- 11) T
8t+11
T (--t)
(107)
where the superscipt denotes a matrix exponential. See Jing et al. (2024b) for precise details.
With both the reference and generated discretized paths, we compute the following metrics:
29

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machiiee Functional
1. Jensen-Shannon Divergence. Consider the probability of each MSM state based on the frequency which it is visited in the discretized paths. We compute these probabilities for both the reference and generated paths, and compute the JSD between the resulting categorical distributions.
2. Transition Negative Log Likelihood. The average negative log likelihood of a transition from one discretized MSM state to the next, averaged over transitions from all generated paths with nonzero probability under the reference MSM. Since there are I- 1 individual transitions for a path of length L, under the Markovian assumption the average negative
T-Z--t1) 1) TT
log likelihood of transition given by - - log ((s = -
log
3. Fraction of Valid Paths. The fraction of generated paths with nonzero probability under the reference MSM.
When considering replicate MD simulations of different lengths (e.g 1 uus, 10 ji8), we fit a MSM to the simulations using the same discretized clusters as were used to fit the reference MSM, and sample 1,000 paths of length L = 20 in the same way described above.
We note that the reported metrics are sensitive the choice of the MSM trajectory length L used for the reference simulations. Varying L changes the time horizon of the transition and yields qualitatively different reference paths. Since we chose L ==200 20, the reported metrics reflcc the correspondence of generated paths with true paths of length 4 ns. We also find that in practice, the optimal choice of At in the OM action (i.e., the choice that yields the best MSM metrics as defined above) results in time horizons T = LAt (where L here is the number of path discretization points during OM optimization) considerably smaller than the reference time horizon of ns (e.g., for the typical choices of = 200, Al = 1fs, we have Tр = =..2ps). other words, OM optimization yields accelerated transition dynamics, because it largely bypasses local/minor fluctuations that would occur in an actual MD simulation.
Committor Function Analysis. The committor function (xx) (see Section 5.1 for definition) is obtainable as the solution to the steady-state backward Kolmogorov equation (BKE) (Hasyim et al., 2022), which is generally infeasible to solve directly numerically for high-dimensional systems. For the fast-folding proteins, we obtain an empirical estimate of the committor function by dividing the TIC configuration space of each protein into 1002 discrete bins, and replacing the expectation in Eq. (101) with an empirical average over trajectories starting from each bin in the reference MD simulations from Lindorff-Larsee et al. (2011). The resulting committor estimates for the fast-folding proteins are shown in Fig. 12.
Model Architecture and Training. Our denoising diffusion and flow matching generative models parameterized by Graph Transformer architecture identical to what was used in Arts et al. (2023) (in the case of diffusion, we use the exact pretrained model from Arts et al. (2023)). To summarize, nodes are featurized by the ordering of each residue in the overall sequence, while edges are featurized by the pairwise C.-Ca, distances. Nodes and edges are then jointly treated as tokens for imput to the Transformer, which updates the token representations. scalar output is obtained by summing learned linear projections of the token representations. Both the denoising diffusion vector field ca and the flow model velocity field vy are parameterized as the gradient of the final scalar output of the model with respect to the input C. coordinates.
For denoising diffusion, we use the pretrained models from Arts al. (2023). For flow matching, train our models. We train models with 3 attention layers for 1 million iterations, using an Adam optimizer with learning rate of 0.0004 and a cosine annealing schedule reducing to minimum learning rate of 0,00001. We use an exponential moving average with a 0.995. The diffusion models use 1,000 integration steps at inference time, while the flow matching models use 10 steps.
Protein-specific training and architecture hyperparameters are given in Table 4.
Table 4. Architecture and training hyperparameters for diffusion and flow matching generative models on fast-folding protcins.
Hyperparameter
Chignolin
Trp-cage
BBA
Villin
Protein G
Batch size
$12
$12
512
512
256
Number hidden features
OM Optimization Details. We list all the optimization hyperparameters used to perform OM optimization on the fast-folding proteins, for both diffusion and flow matching models, in Tables 5 and
30

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Ensaager-Maching Functional
Chignolin
Trp-Cage
0.50
0.25
1.0
-.50
0.8
Committor Probability
-1.00
0.6
BBA
Villin
Protein G
0.2
0.0
0.0
TCL
Figure 12, Empirical committor landscapes for fast-folding proteins. The committor is computed by binning the conformational space into 100 bins measuring the frequency at which reference MD trajectories initiated in each bin reach target state before the start state. The empirical transition ensemble (shown in black) is defined as the level set 0x45..5 < (x)) - Q.55).
Table 5. Hyperparameters used for OM action optimization on fast-folding proteins with diffusion models.
Hyperparameter
Chignolin
Trp-cage
BBA
Villin
Protein G
Number Generated Paths
Action Type
Truncated
Truncated
Full
Truncated
Truncated
Initial Guess Time (Tania))
250
250
250
250
250
Optimization Time (Tqp)
20
20
Optimization Steps
5000
5000
5000
5000
5000
Optimizer
Learning Rate
0.2
0.2
le-5
le-5
le-5
Path Length (L)
200
Action Timestep (drt)
0.001
0.001
0.001
0.005
0.002
Action Friction (t)
Action Diffusivity
Comparison to Transition Paths from Reference MD Simulations. We directly compare our generated paths against the unbiased MD transition trajectories from the D.E.Shaw reference simulations (Lindorff-Larr) Larsen et al., 2011), finding that our paths sample the correct regions of phase space (Fig. 13).
Visualization of Transition Paths. In Figures 14, 15, and 16, provide additional visualizations of transition paths sampled by our diffusion and flow matching models for the fast folding proteins, both in TIC and atomic space.
Transition Paths with Varying Time Horizons We vary the time horizon T, over which OM optimization is performed, and increase the timestep At to maintain the same number of discretization points L for computational efficiency. As shown in Fig. 17, paths with varying time horizons explore different parts of the transition state ensemble (black), with longer horizon time the bottom-right energy well corresponding to the folded state, and shorter paths taking more direct paths between the endpoints.
paths spending
Training without Transition Regions. We provide additional details on the data-starved experiment described in Section 5.3. Transition states are challenging to sample, and therefore may not be abundant reference MD simulations or structural
31

---

Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Table 6. Hyperparameters used for OM action optimization on fast-folding proteins with flow matching models.
Hyperparameter
Chignolin
Trp-cage
Villin
Protein G
Number of Generated Paths
32
Action Type
Truncated
Truncated
Full
Truncated
Truncutedd
Initial Guess Time (Tanial)
7
7
7
7
7
Optimization Time (Tepe)
Optimization Steps
5000
5000
5000
5000
5000
Optimizer
Adam
Adam
SGD
SGD
SGD
Learning Rate
0.2
0.2
lc-4
lc-4
lc-4
Path Length (L)
2000
200
200
200
200
Action Timestep (Att)
0.0008
0.0005
0.001
0.001
0.001
Action Friction (r))
Action Diffusivity (D)
databases, which typically serve as training datasets for generative models. To simulate the scenario in which the underlying dataset is not exhaustive and under-represents the rare, transition regions, we remove 99% of the datapoints for which 0.1 q(x I 0.9, where q(x) is the empirical committor value (described in Committor Function Analysis). Thus, most of the remaining datapoints have committor values close to 0 1, meaning they initiate trajectories which stay in their respective local energy minima without transitioning across the path. For Chignolin, Trp-Cage, and BBA, the subsampling procedure removes 1.4%, 68%, and 85% of the datapoints, respectively. We train diffusion models on these subsampled datasets using the same hyperparameters described earlier, followed by OM optimization between the same endpoints, using the same hyperparameters as used before. As shown in Fig. 18, the produced transition paths are similar to those shown in Figures 4a and 16. The paths still pass through the expected transition state regions (denoted in black), despite having seen them at a much lower frequency during training.
K. Tetrapeptide Experiments
We provide further details on the tetrapeptide experiments from Section 5.4, All experiments were performed on a single NVIDIA RTX A6000 GPU.
Heavy-Atom Representation. Following Jing et al. (2024b), we use a heavy-atom representation of the tetrapeptides (that is, hydrogens are excluded from the representation). The terminal oxygen (OXT) of the C-terminus is also excluded. This results in tetrapeptides with at most 56 atoms each.
Training, We train flow matching model, parameterized by a Graph Transformer with all the same architecture and training hyperparameters used for the fast-folding proteins (Section J), with the only differences being the inclusion of learnable atom type embeddings and the use of padding tokens due to the variable number of atoms in each tetrapeptide. The atom embeddings are concatenated to the residue ordering and the flow timestep to form the node features. We train on dataset of 3109 tetrapeptides simulated in (Jing et al., 2024b), taking 10,000 evenly spaced configurations from the simulations for each tetrapeptide. We use the same optimizer and learning rate scheduler as with the fast folding proteins (J), training for 100,000 steps with batch size of 2048.
OM Optimization on Held-Out Proteins. We perform OM optimization 100 held-out tetrapeptide sequences not seen during training (using the same splits in (Jing al., 2024b)). As in Jing et al. (2024b), we choose endpoints for the transition paths by picking the pair of states from the MSM with the lowest, non-zero probability flux between them, ensuring that the transitions are challenging but feasible to find.
The optimization hyperparameters are given in Table 7.
Energy Minimization. After OM optimization, we add the missing OXT and hydrogen atoms (at neutral pH) back to the
structures using PDBFixer (https:/ / /github. com/ openmm/ f figlfiixee fxer). We then perform up to 200 steps of energy minimization in OpenMM in vacuum with the amber14 force field using L-BFGS. We restrict the Kabsch-aligned RMSD between the initial and final structures Angstrom to ensure that only very minor structural changes occur.
Evaluation. We use the same Markov State Model- evaluation pipeline described in Section for the fast-folding proteins. Following (Jing et al., 2024b), the TIC dimensions are fit on the backbone and sidechain torsion angles. The reported
32

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
Unbiased MD
Sampled Transition Paths
Trajectories
BBA
Villin
Trp-Cage
Figure Sample trajectories from reference, unbiased fast-folding protein MD simulations from D.E.Shaw (left) and paths sampled via OM optimization (right) overlaid on TICA plots with committor values shaded from blue to red. We observe strong agreement in the TICA space, with our OM paths traversing the same transition regions and exhibiting similar structural transitions. While the OM paths are smoother/do not contain the same fluctuations found in unbiased MD trajectories, they capture the same essential features. OM optimization somewhat oversamples transition paths which avoid energy wells. This could be improved by increasing the time horizon of the transition accommodate time spent kinetic traps.
metries are averaged over the 100 held-out test proteins.
Visualization of Sampled Paths. We provide TIC visualizations of sampled transition paths for selected tetrapeptides alongside paths from the reference MSM in Fig. 19.
33

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager- Functional
Chignolin
Trp-Cage
33
BBA
Protein G
Figure 14, Visualization of sampled transition paths from OM optimization with pretrained diffusion models for fast- folding proteins.
Table Hyperparameters used for OM action optimization on tetrapeptides with flow matching models.
Hyperparameter
Value
Number of Generated Paths
16
Action Type
Truncated
Initial Guess Time
7
Optimization Time
0.5
Optimization Steps
250
Optimizer
Adam
Learning
Path Length (E)
100
Action Timestep (At)
0.0001
Action Friction (r)
Action Diffusivity (D)
34

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Chignolin
Trp-Cage
BBA
Protein
Figure 15. Visualization of sampled transition paths from OM optimization with pretrained flow matching models for fast-folding proteins.
Chignolin
Trp Cage
BBA
Villin
Protein
Sampled Transition Paths Transition State Ensemble
Reference Free
Energy Surface
Sampled Transition
Paths (Diffusion Model)
Sampled Transition
Paths
(Flow Matching Model)
Figure 16. Visualization of sampled fast-folding protein transition paths in TIC space.
35

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
T 0.2 ps
T= 0.4 ps
T= 1 ps
T ps
Commitur Prbubiilty
Trp Cage
Sampled Transition Paths
Transition State Ensemble
Figure Varying time horizons path sampling. OM-optimized transition paths for using physical values hyperparameters controlling relative contributions of path and force norm terms in the OM action. Here we set 1ps " and vary the transition time horizon over E (0.2,0.4.1.2)pp and the timestep over A 1..00....002..000,... 0.01)pp ps. which fixes discretization points.
Chignolin
Trp-Cage
BBA
Reference Free Energy Surface
Subsampled Free Energy Surface
Sampled
Transition
Paths
Sampled Transition Paths Transition State Ensemble
Figure 18. Training datasets and sampled transition paths resulting from removing intermediate committor function values. (Top Row) Original training datasets. Middle Row Datasets resulting from removing 99% of datapoints with committor values (obtained empirically) between 0.1 and 0.9. Bottom Transition paths resulting from OM optimization with diffusion model trained the subsampled datasets.
36

---

Action-Minimizationn Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional
Generate Teyerary
Genenned Tnjectrr
Geee sod Vieeaayy
Genernted Tnyertary
AKIR
Generated Vajettrr
Ganamned Trajetsrr
Gonwrated Tajectory
Generated Tejectory
1.5
1.5
1.5
DNTG
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.5
0.5
0.5
0.5
-1.0
3
Ceneutal Trasswyy
Ocnesttl Trgedory
thaudtos
EDOK
0-
-2-
Cenernted Dnjectory
Treseww 1
2njectory
Thuaetory
FDNT
-2
Tnjectary
Gerersted Tejectoy
Tajectay
Genersted Tojeetary 3
IAMI
0-
-1
-1-
-1
-1
-2
-2
-2
-3 -2
1
-3
-2
-1
0
-3
-2
0
1
Figure 19. Sampled transition paths from OM optimization selected, held-out tetrapeptide sequences. The sampled paths diverse and intuitively pass through high density regions in the TIC free energy landscape.
37


Reviews:
[{"review_id": "pJXHVJLHBC", "reviewer": "Reviewer_Z19m", "summary": "The paper introduces Onsager-Machlup (OM) optimization to sample transition paths, claiming three advantages: efficiency, scalability, and flexibility. OM optimization approach produce transition paths in pre-trained generative models, where the core idea is interpreting candidate paths as the denoise-noise SDE allowing tractable computation. Experiments on 2D Muller-Brown potential and fast-folding coarse proteins show that the method produces realistic paths and is also generalized to unseen tetra-peptides.\n\n## Update after rebuttal\n\nI have confirmed the author's response on the scalability and efficiency concern of the proposed method, and raised the score accordingly.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Scalability and efficiency comparison to previous works\n\nTo address the concern about only including results on coarse-grained tetrapeptides, we also present OM optimization results on all-atom tetrapeptides, which contain up to 56 atoms, in [Figure 6r](https://imgur.com/a/naQLWDy). We obtain competitive results on the Markov State Model metrics, similar to coarse-grained tetrapeptides in Figure 4 of the paper. We plan to replace Figure 4 with this all-atom result for the camera-ready version of the paper.\n\nRegarding efficiency relative to past works, MDGen [1] requires a TPS-specific training procedure of about 460 GPU-hours, followed by inference on the order of a few seconds per tetrapeptide path (numbers obtained via direct correspondence with the authors). Meanwhile, we utilize pretrained generative models and perform OM optimization using the learned score function, which requires approximately 30 seconds per tetrapeptide path.\n\nTo account for the increased memory footprint of larger systems, we use a simple path batching technique, which splits up the discretized path into mini batches (e.g split a large, 5,000-point path into mini-batches of size 100). The batches can be optimized either sequentially or trivially parallelized with multiple GPUs at each step of OM optimization, making our method scalable.\n\nAt a broader level, our method is scalable in the sense that we expect continual performance improvements as the underlying atomistic generative models scale up and learn better score functions. This also applies to the growth in size and quality of underlying training datasets (including incorporating experimental data, etc.). We believe that our approach, which requires no TPS-specific training procedure, aligns well with the growing trend of leveraging large-scale, well-tested, general-purpose generative models—a direction already standard in language, speech, and image modeling communities. As high-quality generative models become increasingly available, this synergy will position our method more favorably than existing TPS approaches, which lack such compatibility.\n\n[1] Jing et al. NeurIPS (2024)\n\n> Transition state distribution for 2D Muller Brown\n\nWe show the distribution of sampled transition (i.e highest-energy) states for 2D Muller Brown in [Figure 7r](https://imgur.com/a/CxPwszr). While OM optimization does not capture the full transition state ensemble, this can easily be obtained by initiating MD simulations along the sampled path.\n\n> Diversity metric\n\nWe did not report an explicit diversity metric because to our knowledge, there is no such agreed upon metric in the TPS literature. While not directly a diversity measure, reported distributional distances using the Jensen-Shannon Divergence (JSD) between the MSM state distribution visited by the true and reference transition paths (see Fig. 3d and 4) serve as a proxy for capturing the appropriate path diversity. If the diversity was under- or over-estimated, this would be reflected in the JSD.", "claims_and_evidence": "The paper’s main claim is supported by theoretical arguments (section 3) and experiments (section 5). Specially for the efficiency aspect, transition paths are generated using pre-trained generative models, resulting valid metrics compared to MSM in smaller cost.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria are suited for the problem of transition path sampling, but additional evaluation criteria may strengthen the author’s claim\n\n- 2D Muller brown: qualitative evaluation is given in the main paper, and additional contents in the appendix. A distribution plot of the transition state or the highest energy in 2D Muller brown would clearly show that the OM optimization works well.\n- Fast-folding coarse-grained proteins, tetra-peptides: evaluation based on MSM following MDGen is well done", "theoretical_claims": "I’ve gone through the core theoretical claim, using Onsager-Machlup action for paths under SDE and can be used to produce realistic transition paths with pre-trained generative models.", "experimental_designs_or_analyses": "The experimental design and analyses are concrete and appropriate.\n\n- 2D Muller brown: well-known synthetic testbed for transition path studies\n- Fast-folding protein systems: evaluation and analysis with MSMs\n- Tetra-peptides: transition path generation for unseen data, with evaluation and analysis with MSMs", "relation_to_broader_scientific_literature": "The key contributions of this paper is related to ‘transition path sampling’, where traditional methods have been struggling due to extensive computation. The proposed method seems efficient compared to prior works (however, please check weakness 2).", "essential_references_not_discussed": "It seems all of related works have been comprehensively discussed, with details in Appendix A.", "other_strengths_and_weaknesses": "**Strengths**\n\n1. Originality\n\nCombining OM optimization with pre-trained generative model seems genuinely enough.\n\n2. Presentation\n\nThe paper is well-written, logically structured making it easy to follow!\n\n3. Extensive experiments\n\nAlong with synthetic systems, fast-folding proteins and tetra-peptides experiments are well done and validate the authors claim\n\n\n**Weaknesses**\n\n1. (Minor) diversity for transition paths\n\nWhile “diversity” is not rigorously quantified (no explicit diversity metric is given), the authors do generate multiple paths and report that many are unique and all have non-zero probability under the reference MSM.  \n\n2. Comparison with MDGen [1]\n\nMDGen is a generative model targeting molecular systems with multiple downstream tasks, with transition path sampling being one of them. Like the evaluation in this paper, it also uses MSMs for evaluation. The generalization task for tetra-peptides seems quite similar to that of MDGens. Is there any comparison between MDGen with OM optimization, e.g., OM is efficient to MDGen in terms of GPU hours?\n\n[1] Generative Modeling of Molecular Dynamics Trajectories, NIPS 2024", "other_comments_or_suggestions": "I do not have any other comments", "questions_for_authors": "1. Scalability\n\nThe authors highlighted that the proposed method is advantageous in scalability, proteins and tetra-peptides are done by coarse-graining leading to less than hundred of atoms. Since MDGen[1] models where tetra-peptides are modeled in all atoms, I am confused what the authors imply by scalability. Could the authors provide some details about scalability compared to prior works?\n\n[1] Generative Modeling of Molecular Dynamics Trajectories, NIPS 2024", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "63i06aIsrO", "reviewer": "Reviewer_jnQt", "summary": "The paper proposes a way of using a score-based or a flow based generative model trained to generate molecular configurations to generate transition paths between meta-stable configurations. The paper proposes so be drawing a relation between what would be a limit SDE corresponding to noising and denoising process of DDPM between two adjacent times and the original SDE coming from the physics of the molecular configurations and its energy potential. This relationship allows to adapt the minimization of the Onsager-Machlup functional to the score from the generative model. The paper then proceeds to a numerical evaluation in diverse settings where they compare the proposed approach to state of the art and show that it is able to sample paths that look plausible from several different metrics and comparable with much more compute demanding simulations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Discussion of optimal diffusion time\n\nA couple of ways that $\\tau_\\text{opt}$ is chosen:\n\n1. For diffusion models, a small nonzero $\\tau$ usually works better than $\\tau=0$. Theorem B.1 highlights why, noting $\\bar \\alpha_0 = 1$. Scores closer to the data are weighted lower in a standard DDPM training process. We have run some cosine similarity tests with the true force over Muller-Brown and Alanine Dipeptide to inform the optimal $\\tau_\\text{opt}$ used in these settings (presented in [Figure 8r](https://imgur.com/a/UDyWcxK)) and we used the reported optimal times for fast folders [1].\n\n2. Even if $s_\\theta(x, \\tau_\\text{opt})$ matches the true score well for small enough $\\tau_\\text{opt}$, this vector field can be difficult to optimize over. We have found at times that annealing over the conditioning time $\\tau_\\text{opt}$ helps the OM optimization: start at a value of $\\tau_\\text{opt}$ closer to noise, where $s_\\theta(x, \\tau_\\text{opt})$ is generally a smoother vector field, then gradually over the optimization anneal $\\tau_\\text{opt}$ closer to data. This then motivates us deriving eq. (11) for any latent time, rather than just close to the data distribution.\n\n[1] Arts et al. JCTC (2023)\n\n> “Ad-hoc” nature of the method\n\nPlease see our response to Reviewer bRih, who also raised a similar concern that the hyperparameters controlling the relative contributions of the path, force, and diffusion terms in the OM action are set in an ad-hoc manner. We have refined our method such that these parameters are now directly interpretable as constants with physical units, rather than arbitrary hyperparameters. This makes choosing these hyperparameters intuitive and constrains the hyperparameter search space. More generally, as presented in the Appendix, our method is rigorously derived from well-known action-minimization principles from statistical mechanics, and the equivalence between the learned score function from generative models and the true force field has a clear theoretical justification in the limit of a large, expressive model trained on comprehensive data.\n\n> Erroneous trajectories in flow matching (Figure 7):\n\nThere are a couple factors that make optimized trajectories (even at 0 temperature) look like they’re not following the data distribution:\n\n1. There were fewer plotted samples in order to maintain figure simplicity, but as a result what is not displayed is the fact that the learned distribution does cover a wider band in the transition path, and the converged path is finding the shortest way through this band (hugging the inner wall).\n\n2. The trained flow matching model is not perfect, and has some erroneous local minimum at the center, which gets picked up more at higher temperatures and thus requires a more conservative (lower) $dt$ parameter.\nIf we increase waypoints and are more careful with the optimization, we can obtain more reasonable trajectories even under a potentially erroneous flow matching model, see the now-provided [Figure 9r](https://imgur.com/a/xxpKHrI).\n\n> Using Langevin dynamics instead of the noising-denoising limiting equation\n\nEq. (11) is essentially Langevin dynamics with a learned score at a fixed time-conditioning (note the reasoning works at any time in the noising process, where the vector field is the score of a noised distribution rather than the data distribution). Meanwhile, the standard Langevin dynamics used for sampling/denoising varies the score from $s_\\theta(x, T)$ to $s_\\theta(x, 0)$ throughout the trajectory. \n\nThe motivation for using the fixed-time, noising-denoising limiting equation is the following: our desiderata for the SDE is that the entire trajectory generated by the SDE remains high-likelihood under the data distribution. For variable-time Langevin dynamics, this desiderata is not satisfied, as only the converged/steady-state limit of the dynamics produces data samples. Meanwhile, the noising/denoising equation satisfies this criterion, and the limit process is a valid SDE to which we can apply OM optimization.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "The proposed datasets are pertinent to the problem at hand.", "theoretical_claims": "yes", "experimental_designs_or_analyses": "Yes, the experiments are pertinent and show the value of the proposed method. I still have an issue with the lack of discussion of the optimal diffusion times. I think it would be interesting to see for example in the toy example what is the tradeoff between diffusion time $\\tau_{opt}$ and the distance between the wells. As I see it, having more distant metastables states would lead to an augmentation of the $\\tau_{opt}$ used.", "relation_to_broader_scientific_literature": "The paper proposes a method relying on pre-trained generative models for the molecular configuration for sampling transition paths in a much faster way than the other proposed methods in the literature, such as learning the control drift through reinforcement learning or direct simulation. Although it is has less theoretical guarantees, it can still be a valuable addition.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper present an original idea for using generative models to sample joint transition paths, which open new research venues. The main weakness of the paper is in my opinion the lack of further investigation around the score based model and it's trade-offs and also the fact that the proposed method is rather adhoc as said in the limitations sections.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Could the authors further comment figure 7? It seems that in this case flow matching is actually able to sample configurations around the transition are but the actual OM minimization lead to some erroneous trajectories.\n2. Why not simply use Langevin dynamics with the learned score instead of the noising denoising limiting equation? (11)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "guHvZU43HA", "reviewer": "Reviewer_dGUc", "summary": "The manuscript focuses on transition path sampling (TPS), which involves identifying high-probability paths between two states or points on an energy landscape. The authors combine generative models trained to sample temporally independent states from an energy landscape with the task of transition path sampling. These paths arise from a stochastic differential equation (SDE) constructed from a diffusion or flow-matching model. The authors observe that finding high-likelihood transition paths is equivalent to minimizing the Onsager-Machlup action functional. This connection enables the authors to repurpose pre-trained generative models for TPS. They demonstrate their approach on protein and molecular systems.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Comparison against baselines and incorporation of other generative models\n\nPlease see our response to reviewer bRih, in which we compare our OM optimization approach on alanine dipeptide with two traditional approaches for transition path sampling: Markov Chain Monte Carlo (MCMC)  and metadynamics. We find that OM optimization is several orders of magnitude faster than these approaches and finds transition paths with energy barriers that agree with the reported literature values. We agree that it would also be very interesting to incorporate AlphaFlow and other recent generative models into our OM framework. We have done some preliminary work with AlphaFlow, but found it challenging to work with due to its reliance on pre-trained OpenFold/ESM models, which are extremely memory-intensive. We are actively working to address these challenges in future work.\n\n>Unsuitability of denoising diffusion SDE for transition path sampling\n\nThe full direct denoising process starts with maximum time conditioning ($t=T$) at Gaussian noise (i.e. the target of the learned vector field $\\epsilon_\\theta(x, T)$ is the score of a pure Gaussian). Therefore, optimizing paths to have high likelihood under the denoising SDE would provide a “force field” near the start of the trajectory that just forces it to the origin. Thus the denoising SDE is ill-posed for our setting of TPS, since we want our SDE to produce paths which are high-likelihood under the data distribution at all times, not just at the final denoising step.\n\n>Example of specialized training procedures in prior work and their drawbacks.\n\nSpecialized training procedures in past works include solving a Schrodinger Bridge problem via Stochastic Optimal Control, reinforcement learning, or differentiable simulations (see our submission for citations). These training procedures are expensive because they involve running MD simulations during training, and most trajectories fail to reach the target state, yielding sparse rewards and high cost. A simulation-free approach that guarantees endpoint constraints was introduced in [1]. However, all of these techniques require a training process which is **unique to transition path sampling and must be repeated for every new system of interest (since they do not utilize any training data)**, which is expensive. Meanwhile, our approach indirectly leverages the data and compute used to train atomistic generative models (which can be used for many tasks beyond TPS, notably conformational sampling), using a lightweight test-time OM action minimization procedure.\n\n>Comparison to Doobs Lagrangian\n\nDoobs Lagrangian [1] is a data-free method, meaning it relies only on querying the underlying energy/force function, and requires a bespoke training procedure for every molecular system of interest in order to learn the optimal bias potential for TPS. Meanwhile, our approach leverages pretrained generative models off-the-shelf, with no specialized training procedure for TPS, and can be used across chemical space (barring excessive distribution shift from the generative model’s training data), as we demonstrate in our “Generalization to Unseen Tetrapeptides” experiments. \n\n>Using OM optimization for generative tasks\n\nOptimizing the denoising process directly could be feasible, and indeed falls into the OM optimization framework (essentially by changing $\\tau_\\text{opt}$ between endpoints, similar to $\\tau$ in the denoising process). However, since our paper’s setting focuses on optimizing trajectories that are high likelihood under the data distribution throughout the entire trajectory (rather than just at the endpoint), we have not run such experiments here. This is an interesting idea for future work!\n\n>Zero-shot property of OM optimization.\n\nOM optimization only occurs at test-time, and is a gradient descent procedure over the transition path, not over the generative model weights. The generative model weights are completely fixed throughout our procedure, which is why we call it “zero-shot”.\n\n>What time parameter is used in OM optimization?\n\nDue to space constraints, please see our response to reviewer jnQt, who asked a similar question.\n\n[1] Du et al. NeurIPS (2024)", "claims_and_evidence": "The claims made in the manuscript are backed with theoretical analysis and experiments.", "methods_and_evaluation_criteria": "The method is grounded in strong theoretical results.", "theoretical_claims": "The authors state that the SDE arising from Denoising Diffusion Probabilistic Models (DDPM) is a natural candidate for TPS, but not necessarily a good one. Was this experimentally tested? If so, it may be useful to include a reference to the experiment.\n\n>While the denoising (i.e., sampling) process of a DDPM (see Appendix B.1) may appear to be a natural candidate, a closer inspection reveals that it is unsuitable, as it optimizes for different likelihoods at different points along the trajectory. A large portion of the denoising trajectory thus has low likelihood under the data distribution. Therefore, we need to consider an alternative approach.", "experimental_designs_or_analyses": "The results are interesting and demonstrate strong performance. However, I believe comparisons against existing methods or simpler baselines are lacking. For example, given access to both the score and probability, you could try a simple Markov Chain Monte Carlo (MCMC) algorithm like Metropolis-Adjusted Langevin Algorithm (MALA) or Hamiltonian Monte Carlo (HMC)? Similar works, like AlphaFlow, are relevant to molecular dynamics, and there exist diffusion and flow-matching generative models for proteins that could be integrated within your framework. This suggestion is to enhance comparisons with existing methods; if you know of any others that would be appropriate, that would be welcome.", "relation_to_broader_scientific_literature": "The authors cleverly combine existing theoretical findings and adjust flow matching (FM) to their framework. Their work could be used for various applications that already use diffusion models or flow-matching.", "essential_references_not_discussed": "[1] AlphaFlow could be cited as relevant works.\n\n[1] Jing, B., Berger, B., & Jaakkola, T. (2024). AlphaFold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845.", "other_strengths_and_weaknesses": "I found the introduction and background sections well written and very informative. The method is theoretically grounded, and its explanation is clear. \n\nThe primary drawback is the lack of extensive comparisons in the experimental sections.", "other_comments_or_suggestions": "- > However, these approaches rely on highly specialized training procedures and fail to exploit the growing quantity of atomistic simulation and structural data\n\nIt would be helpful to provide an example of these specialized training procedures or at least elaborate on why this is a drawback of existing approaches.", "questions_for_authors": "- I would like to confirm that the reason it is zero-shot is that the Onsager-Machlup (OM) minimization only occurs at inference? There is no need for a fine-tuning stage?\nRepurpose pre-trained generative models for TPS in a zero-shot fashion.\n\n- Have you attempted using the OM for generative tasks in a manner similar to your current approach, but starting with a noise representation?\n\n- It would be beneficial to clarify the differences between your work and [2].\n\nIt is mentioned in the limitations, but what are the advantages of your method compared to this one (if I understand correctly, they need to train or fine-tune)?\n\n[2] Du, Y., Plainer, M., Brekelmans, R., Duan, C., NoÅLe, F., Gomes, C. P., Aspuru-Guzik, A., and Neklyudov, K. Doob’s Lagrangian: A sample-efficient variational approach to transition path sampling. arXiv preprint arXiv:2410.07974, 2024.\n\n- In diffusion and flow-matching models, the score of the vector field is usually parameterized by a time variable that interpolates between noise and data. Since, in your case, this initial point is not necessarily noise, which time parameter do you use? I assume the time variable from Eq. 11 will not be used to compute the score $s_\\theta(x,\\tau)$, especially since $\\tau$ needs to fall within the interval originally used to train the model.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AnWv3eaU0e", "reviewer": "Reviewer_bRih", "summary": "This paper presents a new method for transition path sampling in molecular systems by combining generative models with the Onsager-Machlup action functional. The authors show how pre-trained generative models (specifically denoising diffusion and flow matching) can be repurposed to find high-probability transition paths between stable configurations without requiring specialized training. The author's approach interprets candidate paths as trajectories sampled from stochastic dynamics induced by the learned score function, making finding transition paths equivalent to minimizing the Onsager-Machlup action. The method is demonstrated on a 2D Müller-Brown potential, fast-folding proteins, and tetrapeptide sequences, showing it can generate physically realistic transition pathways. Compared to traditional molecular dynamics simulations, their approach produces higher-quality transition paths with significantly less computational cost.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">Comparison with enhanced sampling baselines\n\nTo our knowledge, there’s no widely accepted force field for the $\\alpha$-Carbon coarse-graining used for the fast-folding proteins, so benchmarks are challenging. We now benchmark on all-atom alanine dipeptide, a standard test system for TPS. We compare with MCMC (shooting) and metadynamics. See [Table 1r](https://imgur.com/a/iPmIUbp) for accuracy and efficiency results, and [Figure 1r](https://imgur.com/a/SImqo3n) for sampled transition paths.\n\nOM optimization is considerably faster than the baselines. Internal energy barriers were about 10-12 kcal/mol, in line with literature. \n>Quantitative evaluation of transition paths (reaction rate, committor function, free energy profile)\n\nThese are known to be challenging to compute accurately for high-dimensional systems. With this in mind, we provide results to show how our OM optimized paths can be used to compute these quantities. We’re happy to include more in-depth results for the camera-ready version. \n\nFree Energy: Our OM-optimized paths can serve as a guide for umbrella sampling, from which we can calculate free energy profiles. We do this for Alanine Dipeptide and achieve the following free energy profile, which closely approximates the barrier of 6-8 kcal/mol obtained from metadynamics: [Figure 2r](https://imgur.com/a/3ci9dXh)\n\nCommittor Function/Transition Rates: The Backward Kolmogorov Equation (BKE) provides a principled algorithm [1] to estimate the committor function and transition rates using the paths from OM optimization. See [Figure 3r](https://imgur.com/a/BhLLOKs) for results on the 2D Muller-Brown potential, where we achieve accurate committor and rate estimation, with a margin of error comparable with or better than past works. Note that it is often challenging to even compute the reaction rate within the correct order of magnitude [1,2].\n>Guidance on setting hyperparameters \n\nWe now set hyperparameters directly as values with physical units: for example setting $\\gamma$ and $D$ to be friction and diffusion coefficients, and $\\Delta t$, $L$ from the time horizon and desired fidelity. We provide an example of new paths with varying time horizons in [Figure 4r](https://imgur.com/a/s9DhO3R). \n>Markov State Models (MSM) do not confirm dynamical accuracy\n\nOur reference MSMs are fit on long D.E.Shaw simulations which extensively sample the Boltzmann distribution, including many transition events. Therefore, the reference MSMs are a useful low-dimensional approximation of the underlying dynamics of the unbiased MD simulation, which is also standard practice in past works [3,4].\nWe have also directly compared our generated paths against the unbiased MD transition trajectories from D.E.Shaw: [Figure 5r](https://imgur.com/a/7m7Npne)\n>OM Optimization on chemical reactions\n\nWe conducted initial experiments using a machine learning force field (MLFF) trained on the Transition-1x dataset. OM interpolation finds transition states on average 1.4eV above the reference transition state (found using Nudged Elastic Band (NEB) calculations with DFT), significantly better than running NEB calculations with the MLFF (1.91eV above) and uses 10x fewer force evaluations. A comprehensive analysis would be a paper in and of itself, so we consider this an interesting future direction.\n>Zero-shot generalization on unseen tetrapeptides\n\nWhile the force field is the same for all tetrapeptides (see [4]), it is sequence-dependent: the energy landscape is different for held-out proteins compared with training proteins. Thus, our results in Section 5.3 effectively demonstrate generalization to new interaction potentials.\n>Rare events underrepresented in training data\n\nWe agree that this is an inherent challenge associated with data-driven methods. However, we have demonstrated experimentally in Section 5.2 that OM optimization is resilient to a degree of data sparsity and still identifies realistic transition paths. As we scale up to larger models like BioEmu [5], we expect resilience to sparsity to get stronger.\n>Preferring low-energy interpolations\n\nWe can still achieve diversity in sampled transition paths by leveraging stochasticity of the generative model to get diverse initial guesses (Alg. 2, Sec. G) – converged paths thus tend to find diverse local minima. Further, $\\tau_\\text{initial}$ can be tuned to vary the diversity of the initial guesses and subsequently produced paths.\n>Solving CV selection problem\n\nWe don’t claim to solve the CV selection problem. We just note that our method does not require CV to operate.\n>OM Optimization on completely different force fields or chemical environments\n\nThis is an important area for future work, and we believe that using more performant models such as BioEmu [5] as the underlying score estimator would aid in achieving this.\n[1] Hasyim et al. JCP (2022)\n[2] Rotskoff, et al. PMLR (2022)\n[3] Arts et al. JCTC (2023)\n[4] Jing et al. NeurIPS (2024)\n[5] Lewis et al. arxiv: 2024.12.05.626885", "claims_and_evidence": "The overall claims focus on the efficiency and physically realistic transition path generated based on the OM algorithm, supported by its implementation on systems such as Müller-Brown potential and coarse-grained proteins.\n\nComparisons are made between OM-generated paths and MD simulations in a few systems with wall-clock time comparison presented in Figure 3c shows that OM optimization is faster than brute-force MD. However, the comparison lacks benchmarking against enhanced sampling methods (e.g., metadynamics, umbrella sampling, weighted ensemble, transition interface sampling). These are the actual competitors, not brute-force MD. Therefore, The computational cost of optimizing the OM action versus performing biased MD simulations is not clearly reported.\n\nRegarding the generated transition path, the claim of its physical realism requires additional proof. The paper claims that minimizing the OM action produces physically meaningful transition paths, even for systems unseen during training, given the comparison between transition path and free energy landscapes. Since the generative models are trained on equilibrium data, the core issue is that the score-based generative models do not inherently capture rare transition pathways. The paths generated may be interpolations that resemble plausible transitions rather than actual dynamical pathways governed by the correct physics (figure 13 and 14). Besides, MSM-based validation is useful but does not confirm the dynamical accuracy of the generated paths. The MSM itself is only as good as the data used to construct it.\n\nAnother claim that requires further examination is the zero-shot generalization statement. However results on held-out tetrapeptides suggest that transition paths can be generated for new sequences, different tetrapeptides have unique energy landscapes, and it is unclear how well a generative model trained on one set of peptides generalizes to another with different interaction potentials. Generative models trained on equilibrium data are likely to underrepresent high-energy, transition-state conformations, which means the model might struggle with generating physically accurate rare event dynamics.\n\nThe strong claim of solving the CV selection problem is not solid enough -- the method optimizes transition paths based on the OM action but does not guarantee exploration of the full transition ensemble. If the generative model was trained on biased datasets (e.g., experimental structures, short MD trajectories, not-complete stationary structures), then the inferred transition paths will be similarly biased.", "methods_and_evaluation_criteria": "The main method introduced in the paper is OM action minimization using generative models (specifically denoising diffusion models and flow matching models) to infer likely transition pathways between molecular configurations. The key steps include training a generative model on equilibrium molecular conformations; using the learned score function as an approximate force field and optimizing paths between two states by minimizing the OM action.\n\nThe OM action is based on path probability, but most probable does not equal dynamically correct. The method may find smooth interpolations rather than true dynamical pathways dictated by physical kinetics. Besides, there is no comparison is made with unbiased MD simulations to check if the generated paths follow the correct transition state kinetics.\n\nAnother issue is that the optimization heavily depends on parameters controlling diffusion, drift, and friction, but the method does not provide a systematic way to tune these parameters for different molecular systems. For example, the experiments show that lower diffusivity (Figure 2) gives a better estimation of the intrinsic reaction coordinate (minimum energy pathway that connects the saddle point and two minima) while the larger diffusivity is way off the accurate transition path. \n\nA more robust evaluation of TPS methods should include:\nComparison with well-established TPS methods (e.g., umbrella sampling, metadynamics, transition path sampling).\nValidation against unbiased MD simulations to assess kinetic accuracy.\nQuantitative measures of transition rate, free energy profile, and committor function accuracy.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The issue has been described in the previous question,  the evaluation experiments need stronger validation. The proposed method is innovative, but its evaluation is insufficient to justify its claims of accuracy, efficiency, and generalizability.", "relation_to_broader_scientific_literature": "The contribution of the paper may benefit the study of transition paths but still lacks confidence due to incomplete validation, limited benchmarking, and potential oversimplification of transition path dynamics. While the idea of leveraging generative models for transition path sampling (TPS) is innovative, the connection between OM action minimization and physically meaningful transition paths remains insufficiently supported. Below, I examine how the paper relates to the broader scientific literature, highlighting where it builds upon previous work and where gaps remain. Due to the high degree of freedom for reaction systems (DOE = 3N-5), the complexity of the potential energy surface can not be solved in this paper. The unique characteristics of the transition path connecting the transition state— which appears as a saddle point in the reaction coordinate projection but as a minimum in other projections—cannot be predicted solely based on the minima.", "essential_references_not_discussed": "The related works are properly cited.", "other_strengths_and_weaknesses": "Strength: The paper presents a scalable, and computationally efficient approach to transition path sampling using generative models and OM minimization. \n\nWeakness: The paper lacks crucial validation against enhanced sampling methods and transition rate calculations, leaving uncertainty about its physical accuracy. If these issues are addressed in future work, the approach could impact molecular simulations and accelerate the discovery of reaction pathways.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How does OM optimization compare to well-established enhanced sampling techniques in terms of both efficiency and accuracy? Without this comparison, the claim of “higher efficiency” is weak.\n\n2. How do these transition paths compare quantitatively with TPS or enhanced sampling methods in terms of transition rates and free energy barriers?\n\n3. How does OM optimization perform on systems where the generative model was trained on an entirely different force field or chemical environment? Does it still yield reasonable transition paths?\n\n4. Does this method capture all relevant transition pathways, or does it preferentially generate low-energy interpolations that may miss rare but important transitions?\n\n5. How does OM optimization perform with chemical reaction systems (when bond cleavage/formation is happening)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The author present an application of the Onsager-Machlup action minimization functional to sample transition paths in pretrained molecular ensemble generation models, e.g. pre-trained Boltzmann Generators or emulators. While the approach is conceptually simple and well-known in the molecular simulation community (work by Vanden-Eijden and others is discussed in the paper), however, their practicality have been limited by computational demands in its implementation. The application to pre-trained generative models is as far the reviewers and I am concerned, new. However, when applied this setting the sampled transition paths will only be as accurate as the underlying surrogate allows. The reviewers are positive overall, however, their enthusiasm appears to hinge to some degree of the promise of scaling these approaches to larger and more challenging systems than the ones explored the in the paper. The authors acknowledged this and promises a more extensive set of systems in their final paper.
Decision: Accept (poster)