[Paper Review] Differentiable Simulations for Enhanced Sampling of Rare Events

Martin Šípka, Johannes C. B. Dietschreit+2 more|arXiv (Cornell University)|Jan 9, 2023

Machine Learning in Materials Science4 citations

TL;DR

This paper introduces differentiable simulations (DiffSim) to discover and enhance sampling of rare chemical events without relying on preselected collective variables (CVs). By formulating reaction path discovery and biasing potential learning as a joint, end-to-end optimization via path-integral loss, the method enables stable, efficient training through partial backpropagation and graph mini-batching, successfully identifying transition paths in the Muller-Brown potential and alanine dipeptide without prior CV knowledge.

ABSTRACT

Simulating rare events, such as the transformation of a reactant into a product in a chemical reaction typically requires enhanced sampling techniques that rely on heuristically chosen collective variables (CVs). We propose using differentiable simulations (DiffSim) for the discovery and enhanced sampling of chemical transformations without a need to resort to preselected CVs, using only a distance metric. Reaction path discovery and estimation of the biasing potential that enhances the sampling are merged into a single end-to-end problem that is solved by path-integral optimization. This is achieved by introducing multiple improvements over standard DiffSim such as partial backpropagation and graph mini-batching making DiffSim training stable and efficient. The potential of DiffSim is demonstrated in the successful discovery of transition paths for the Muller-Brown model potential as well as a benchmark chemical system - alanine dipeptide.

Motivation & Objective

• Address the challenge of discovering reaction pathways in high-dimensional molecular systems where traditional enhanced sampling relies on manually chosen collective variables (CVs).
• Overcome the 'chicken-and-egg' problem in CV selection, where identifying CVs requires prior knowledge of the reaction path.
• Develop a robust, differentiable simulation framework capable of learning biasing potentials that enhance rare event sampling without requiring expert-defined CVs.
• Enable end-to-end optimization of both the reaction path and the biasing potential using a unified path-integral loss function.
• Demonstrate the method’s effectiveness on benchmark systems, including the Muller-Brown potential and alanine dipeptide, with minimal prior assumptions.

Proposed method

• Formulate a path-integral loss function that measures the likelihood of a trajectory crossing an energy barrier, serving as a differentiable proxy for rare event occurrence.
• Train a learnable biasing potential using backpropagation through the entire simulation trajectory, enabling end-to-end optimization of the bias to enhance sampling.
• Implement partial backpropagation to stabilize gradients and prevent vanishing/exploding gradients in long trajectories.
• Introduce graph mini-batching to enable stochastic gradient descent, improving training stability and convergence.
• Use Langevin dynamics to introduce a finite memory horizon, which naturally suppresses adjoint growth and stabilizes the optimization process.
• Apply the .detach() operation to prune the computational graph and suppress high-frequency oscillations in gradients.

Figure 1 : Comparison of the adjoint evolution for original and partially detached graphs for simulations on the 2D Müller-Brown potential with parameters given in Appendix F .

Experimental results

Research questions

• RQ1 Can differentiable simulations be used to discover reaction pathways without relying on preselected collective variables?
• RQ2 Can a unified, end-to-end optimization framework effectively learn both the reaction path and the biasing potential in a single training process?
• RQ3 How stable and efficient is the training of differentiable simulations when applied to high-dimensional molecular systems with complex energy landscapes?
• RQ4 Can the method successfully enhance sampling of rare transitions in realistic chemical systems like alanine dipeptide without prior knowledge of key dihedral coordinates?
• RQ5 To what extent can the learned biasing potential reproduce known free energy profiles, such as the PMF from metadynamics?

Key findings

• The method successfully discovered the transition path in the 2D Muller-Brown potential without any prior knowledge of the collective variables, achieving convergence even in the presence of noise.
• In the 5D noisy version of the Muller-Brown potential, the method still achieved a high probability of observing transitions, demonstrating robustness to irrelevant degrees of freedom.
• For alanine dipeptide, over 80% of trajectories completed a transition within 10 ps after training, indicating effective enhancement of rare events.
• The average bias potential learned by the method closely matched the reference PMF obtained from metadynamics, validating its accuracy in capturing the free energy landscape.
• The method enabled the construction of a reduced representation (CV) from biased trajectories, showing that the learned bias implicitly identifies relevant reaction coordinates.
• The use of partial backpropagation and graph mini-batching significantly stabilized training, enabling reliable optimization of the differentiable simulation.