Learning Collective Variables for Enhanced Sampling from BioEmu with Time-Lagged Generation

TL;DR

The paper tackles the slow-dynamics identification bottleneck in molecular dynamics by proposing BioEmu-CV, a framework that learns a slow CV from a frozen BioEmu foundation model through time-lagged generation conditioning. By freezing BioEmu and training a lightweight encoder with a time-lagged score-matching objective, BioEmu-CV captures slow macroscopic dynamics while ignoring fast fluctuations. The authors validate the approach on three fast-folding proteins using two downstream tasks—on-the-fly OPES for free energy estimation and CV-steered MD for transition-path sampling—showing competitive accuracy, improved state discrimination, and interpretable CVs that align with known structural features. This work also establishes a new benchmark for machine-learned CVs on relatively large biomolecular systems, providing a scalable path toward automated CV discovery in enhanced sampling.

Abstract

Molecular dynamics is crucial for understanding molecular systems but its applicability is often limited by the vast timescales of rare events like protein folding. Enhanced sampling techniques overcome this by accelerating the simulation along key reaction pathways, which are defined by collective variables (CVs). However, identifying effective CVs that capture the slow, macroscopic dynamics of a system remains a major bottleneck. This work proposes a novel framework coined BioEmu-CV that learns these essential CVs automatically from BioEmu, a recently proposed foundation model for generating protein equilibrium samples. In particular, we re-purpose BioEmu to learn time-lagged generation conditioned on the learned CV, i.e., predict the distribution of molecular states after a certain amount of time. This training process promotes the CV to encode only the slow, long-term information while disregarding fast, random fluctuations. We validate our learned CV on fast-folding proteins with two key applications: (1) estimating free energy differences using on-the-fly probability enhanced sampling and (2) sampling transition paths with steered molecular dynamics. Our empirical study also serves as a new systematic and comprehensive benchmark for MLCVs on fast-folding proteins larger than Alanine Dipeptide.

Figures (19)

• Figure 1: Overview of our framework and evaluation simulation.(Left) We train an encoder on top of the conditions of a frozen molecular foundation model to learn collective variables (CVs) for protein, highlighted in red dotted lines, on top of a frozen pre-trained BioEmu. (Right) Two downstream tasks for the slow degree of freedom, free energy estimation and transition path sampling.
• Figure 2: Free energy (top) and PMF (bottom) estimation from 1 $\mu s$ OPES simulations for Chignolin. We average over four seeds. Green dotted lines indicate the reference value, and blue lines refer to free energy difference during the OPES simulations. Solid lines refer to the mean, and shaded areas are the standard deviation. DeepTICA shows negative values beyond -1 in PMF, even normalized to the DESRES trajectory, falls short of accurately capturing the reference PMF.
• Figure 3: 3D Visualization of transition paths. The sampled folding pathways of Chignolin, Trp-cage, and BBA by steered MD with BioEmu-CV. For simplicity, we visualize the $C_{\alpha}$ coordinates.
• Figure 4: Transition paths projected onto TICA coordinates, sampled by MLCV steered MD. The white star and circle each refer to the representative unfolded and folded state, with each paths colored differently. The red convex hull is the folded states defined as RMSD from the representative folded state with cut off 2, 2, and 2.5 Å for Chignolin, Trp-cage, and BBA, respectively.
• Figure 5: MLCV sensitivity to $C_\alpha$-wise distances. We plot the top $C_\alpha$-wise distances with the highest sensitivity for each MLCV, where the $x$ and $y$ axis each denotes the residue index for input distances and the corresponding sensitivity value. For each sensitivity plot, we visualize top sensitive distances in the unfolded and folded state, with colors highlighted in the sensitivity plot.