Contrastive Regularization of Machine Learning Potentials

Technical Summary: Contrastive Regularization of Machine Learning Potentials

Problem Statement
Machine Learning Interatomic Potentials (MLIPs) are typically trained as regressors to predict energies and forces from Density Functional Theory (DFT) data, minimizing Mean Squared Error (MSE). While these models often achieve chemical accuracy on held-out test sets, the authors demonstrate that they frequently fail as generators when deployed to drive molecular dynamics or Monte Carlo simulations.

The core issue is statistical rather than architectural. In high-dimensional configuration spaces, a finite training set covers a vanishingly small fraction of the space. Standard MSE minimization constrains the potential only where data exists, leaving the rest of the configuration space unconstrained. Consequently, expressive neural networks inevitably carve out spurious low-energy minima in these unvisited regions. When used as samplers, trajectories drift into these "ghost" minima, resulting in unphysical configurations and thermodynamic observables (e.g., free-energy profiles, structural distributions) that deviate sharply from the reference, even though pointwise energy and force predictions remain accurate.

Methodology: Contrastive Regularized MSE (CRMSE)
To address this, the authors introduce Contrastive Regularized MSE (CRMSE), a post-training procedure that augments the standard MSE loss with a contrastive term derived from the Kullback–Leibler (KL) divergence between the model's implicit Boltzmann distribution and the target distribution.

Key methodological components include:

• Objective Function: The total loss is $L_{CRMSE} = L_{MSE} + \lambda L_{CR}$. The MSE term maintains accuracy on the training data, while the contrastive term ($L_{CR}$) optimizes the probability distribution.
• Contrastive Mechanism: The contrastive term is derived from minimizing the KL divergence. It acts as a variational free energy objective: it lowers the energy of "positive" samples (the training data) and raises the energy of "negative" samples (configurations drawn from the model's current distribution).
• Self-Consistent Sampling: Since the model's distribution is intractable, negative samples are generated using the Unadjusted Langevin Algorithm (ULA). The authors employ Persistent Contrastive Divergence (PCD), where a replay buffer of negative configurations is maintained and updated with a few ULA steps per iteration.
• Post-Training Nature: CRMSE is applied as a post-training step to an already MSE-optimized potential. It requires no new ab initio data and no architectural changes. The model effectively "labels its own mistakes" by identifying configurations it drifts into during simulation and raising their energy to create a protective barrier around the physical basin.

Key Results
The method was validated on the ethanol and aspirin molecules from the MD17 dataset using a Graph Neural Network with Local Frame (GNN-LF) architecture.

1. Failure of Standard MSE: Standard MSE-trained models achieved high accuracy on held-out energies and forces. However, when used in Metropolis-adjusted Langevin algorithm (MALA) simulations, trajectories rapidly escaped the physical basin, settling in spurious minima tens of kcal/mol below the reference energy. This resulted in unphysical interatomic distances and incorrect energy distributions.
2. Correction via CRMSE:
   • Sampling Stability: CRMSE post-training confined the sampler to the physical basin. The sampled energy distributions and interatomic distance histograms recovered near-quantitative agreement with DFT references.
   • Thermodynamic Observables: Crucially, CRMSE recovered the free-energy profiles along dihedral angles (H-C-C-O for ethanol and ester dihedral for aspirin) to near-quantitative agreement with DFT. It corrected the underestimation of torsional barriers and collapsed the inter-chain dispersion (variance between independent simulation chains) to within statistical noise.
   • Preservation of Accuracy: The correction did not degrade the model's predictive accuracy on the training manifold; energy and force errors on held-out data remained within chemical accuracy.
   • Robustness: The method remained effective even when the training set was sharply reduced from 950 to 200 configurations, demonstrating improved data efficiency and robustness to data scarcity.

Significance and Claims
The paper claims that reliable sampling depends less on the volume of training data than on training the model against the distribution it produces. The authors argue that distribution-level training is a distinct requirement, not merely a refinement of regression accuracy.

• Distinct Requirement: The work highlights that pointwise accuracy does not guarantee a correct sampling distribution. A potential meant to be sampled must be trained to reproduce its own sampling distribution.
• Complementarity: CRMSE is presented as complementary to active learning. While active learning injects new ab initio data to improve accuracy in off-manifold regions, CRMSE restores the correct support of the sampling distribution without new data, suppressing spurious minima by raising their energy.
• Generality: The authors suggest this lesson applies broadly to AI settings where learned regressors are deployed as generators. The method is particularly significant because it repairs a potential that is already "well-optimized" by standard regression metrics, using only a cheap post-training step and no additional reference calculations.

The authors conclude that for surrogate models used as generators, the training objective must account for the full probability density defined by the potential, a requirement structurally omitted by standard MSE minimization.