[Paper Review] Enhanced sampling of robust molecular datasets with uncertainty-based collective variables

Aik Rui Tan, Johannes C. B. Dietschreit+1 more|arXiv (Cornell University)|Feb 6, 2024

Computational Drug Discovery Methods4 citations

TL;DR

This paper proposes using uncertainty from a single machine learning model as a collective variable (CV) to guide enhanced sampling in molecular dynamics, enabling efficient exploration of rare and high-energy configurations. By biasing simulations toward regions of high prediction uncertainty, the method achieves broader coverage of configuration space—demonstrated on alanine dipeptide—improving data set diversity and reducing extrapolative errors in machine learned interatomic potentials (MLIPs).

ABSTRACT

Generating a data set that is representative of the accessible configuration space of a molecular system is crucial for the robustness of machine learned interatomic potentials (MLIP). However, the complexity of molecular systems, characterized by intricate potential energy surfaces (PESs) with numerous local minima and energy barriers, presents a significant challenge. Traditional methods of data generation, such as random sampling or exhaustive exploration, are either intractable or may not capture rare, but highly informative configurations. In this study, we propose a method that leverages uncertainty as the collective variable (CV) to guide the acquisition of chemically-relevant data points, focusing on regions of the configuration space where ML model predictions are most uncertain. This approach employs a Gaussian Mixture Model-based uncertainty metric from a single model as the CV for biased molecular dynamics simulations. The effectiveness of our approach in overcoming energy barriers and exploring unseen energy minima, thereby enhancing the data set in an active learning framework, is demonstrated on the alanine dipeptide benchmark system.

Motivation & Objective

• Address the challenge of sparse and unrepresentative training data in machine learned interatomic potentials (MLIPs), which limits model generalizability.
• Overcome the limitation of traditional sampling methods that favor low-energy minima and neglect rare, high-energy configurations critical for robust MLIP training.
• Develop a method that actively targets configuration space regions where MLIP predictions are most uncertain, thereby improving data set coverage and model reliability.
• Eliminate the need for pre-defined, system-specific reaction coordinates by using model uncertainty as a universal, adaptive CV.
• Reduce computational cost by leveraging single-model uncertainty instead of ensemble-based uncertainty estimation, while maintaining effective exploration.

Proposed method

• Use a Gaussian Mixture Model (GMM) to estimate uncertainty from predictions of a single trained neural network interatomic potential (NNIP), treating this uncertainty as the collective variable (CV).
• Apply extended-system adaptive biasing force (eABF) with Gaussian-accelerated molecular dynamics (GaMD) to perform biased molecular dynamics simulations using the uncertainty CV.
• Dynamically adjust biasing strength (γ) and uncertainty threshold (u_cutoff) across active learning iterations to balance exploration and exploitation.
• Iteratively collect new configurations from high-uncertainty regions and retrain the NNIP, progressively improving model accuracy and coverage.
• Calibrate uncertainty estimates using conformal prediction to ensure reliable error bounds on predictions.
• Compare the uncertainty-as-CV approach with uncertainty-as-biasing-energy methods, evaluating both configuration space coverage and prediction error reduction.

Figure 1: (a) , Structure of the alanine dipeptide molecule with carbon (C), nitrogen (N), oxygen (O), and hydrogen (H) atoms labeled in grey, blue, red, and white, respectively. Four backbone dihedral angles $\phi$ , $\psi$ , $\omega_{1}$ , and $\omega_{2}$ are annotated. (b) , Potential mean force

Experimental results

Research questions

• RQ1 Can uncertainty from a single trained MLIP serve as an effective, universal collective variable for guiding enhanced sampling in molecular systems?
• RQ2 Does using uncertainty as a CV lead to better coverage of configuration space, particularly in rare or high-energy regions, compared to traditional sampling or uncertainty-based biasing?
• RQ3 How does the choice of biasing strength (γ) affect the balance between exploration of novel configurations and avoidance of unphysical distortions?
• RQ4 To what extent does the uncertainty-as-CV method improve the generalization and extrapolative accuracy of MLIPs across different regions of the potential energy surface?
• RQ5 Can this approach reduce the need for ensemble-based uncertainty estimation while maintaining effective active learning for data set enrichment?

Key findings

• A biasing strength of γ = 0.005 yielded the most effective exploration, achieving higher predicted uncertainties and better coverage than γ = 0.01, suggesting that lower γ enables more gradual, barrier-aware sampling.
• The uncertainty-as-CV method successfully accessed regions of the φ–ψ dihedral space with φ < 0°, which were poorly sampled in unbiased simulations and less explored in uncertainty-as-biasing-energy approaches.
• Despite stagnating coverage fractions beyond generation 4 in active learning, mean absolute error (MAE) in predicted energy and forces continued to decrease, indicating effective enrichment of the model’s extrapolative capability.
• The method achieved effective exploration of non-planar ω₁–ω₂ dihedral angles, a region previously underexplored in similar eABF-GaMD simulations, demonstrating improved sampling of complex torsional motions.
• The uncertainty-as-CV approach outperformed uncertainty-as-biasing-energy in terms of configuration space diversity and model robustness, with no premature simulation termination due to unphysical configurations.
• Single-model uncertainty estimation proved sufficient for guiding effective sampling, reducing computational cost compared to ensemble-based uncertainty methods without sacrificing performance.

Figure 2: Cumulative exploration of configuration space projected onto the $\phi$ - $\psi$ (left column) and $\omega_{1}$ - $\omega_{2}$ (right column) plane of 10 NVT simulations at 300 K. Top) No biasing of any sort, middle) uncertainty-guided eABF, and bottom) uncertainty-guided eABF-GaMD.