Latent Spaces for Langevin Dynamics
Andy Bruce, Alexander Aghili, Razvan Marinescu, Daniel Sabo
TL;DR
The paper addresses the restriction of traditional coarse-graining to quadratically coupled momentum by deriving a general family of embeddings $f$ that render Langevin dynamics thermodynamically exact on non-geometric CG spaces. It shows the coarse-grained free energy can be written as $F(\vec{Q},\vec{P}) = \frac{1}{2}\vec{P}^\top R^{-1}(\vec{Q}) \vec{P} + V(\vec{Q})$, with $R^{-1}(f(\vec{q})) = J_f(\vec{q}) M^{-1}(\vec{q}) J_f^\top(\vec{q})$, ensuring consistency across microstates. The framework enables Langevin sampling directly in latent CG coordinates, separating momentum contributions (computable analytically) from the potential learned by neural networks, and supports hyperparameter tuning via thermodynamic quantities. Practically, this expands the use of Langevin dynamics to ML-derived coarse-grained representations, potentially improving efficiency and accuracy in molecular simulations. Key demonstrations include distance-latent-space encodings and force-matching training of the potential in a neural CG model.
Abstract
In the field of machine learning coarse-grained potentials in molecular dynamics, many propagators require that the effective Hamiltonian is quadratic in momentum, thus limiting the family of coarse-graining functions. In this paper, we derive a general family of coarse-graining embedding functions for which Langevin dynamics samples correctly. These equations have significant implications for molecular simulations and pave the way for Langevin dynamics on non-geometric coarse-graining representations, such as those provided by principal components of component analysis or latent embeddings of molecules obtained from neural networks.