Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional

Comparison with enhanced sampling baselines

To 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 for accuracy and efficiency results, and Figure 1r for sampled transition paths.

OM optimization is considerably faster than the baselines. Internal energy barriers were about 10-12 kcal/mol, in line with literature.

Quantitative evaluation of transition paths (reaction rate, committor function, free energy profile)

These 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.

Free 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

Committor 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 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].

Guidance on setting hyperparameters

We 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.

Markov State Models (MSM) do not confirm dynamical accuracy

Our 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]. We have also directly compared our generated paths against the unbiased MD transition trajectories from D.E.Shaw: Figure 5r

OM Optimization on chemical reactions

We 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.

Zero-shot generalization on unseen tetrapeptides

While 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.

Rare events underrepresented in training data

We 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.

Preferring low-energy interpolations

We 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.

Solving CV selection problem

We don’t claim to solve the CV selection problem. We just note that our method does not require CV to operate.

OM Optimization on completely different force fields or chemical environments

This 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. [1] Hasyim et al. JCP (2022) [2] Rotskoff, et al. PMLR (2022) [3] Arts et al. JCTC (2023) [4] Jing et al. NeurIPS (2024) [5] Lewis et al. arxiv: 2024.12.05.626885

Comparison against baselines and incorporation of other generative models

Please 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.

Unsuitability of denoising diffusion SDE for transition path sampling

The 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.

Example of specialized training procedures in prior work and their drawbacks.

Specialized 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.

Comparison to Doobs Lagrangian

Doobs 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.

Using OM optimization for generative tasks

Optimizing 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!

Zero-shot property of OM optimization.

OM 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”.

What time parameter is used in OM optimization?

Due to space constraints, please see our response to reviewer jnQt, who asked a similar question.

[1] Du et al. NeurIPS (2024)

Discussion of optimal diffusion time

A couple of ways that $\tau_\text{opt}$ is chosen:

1. 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) and we used the reported optimal times for fast folders [1].

2. 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.

[1] Arts et al. JCTC (2023)

“Ad-hoc” nature of the method

Please 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.

Erroneous trajectories in flow matching (Figure 7):

There are a couple factors that make optimized trajectories (even at 0 temperature) look like they’re not following the data distribution:

1. 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).

2. 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. If 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.

Using Langevin dynamics instead of the noising-denoising limiting equation

Eq. (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.

The 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.

Scalability and efficiency comparison to previous works

To 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. 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.

Regarding 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.

To 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.

At 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.

[1] Jing et al. NeurIPS (2024)

Transition state distribution for 2D Muller Brown

We show the distribution of sampled transition (i.e highest-energy) states for 2D Muller Brown in Figure 7r. While OM optimization does not capture the full transition state ensemble, this can easily be obtained by initiating MD simulations along the sampled path.

Diversity metric

We 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.