[Paper Review] Thermodynamically Informed Multimodal Learning of High-Dimensional Free Energy Models in Molecular Coarse Graining

Blake Duschatko, Xiang Fu+5 more|arXiv (Cornell University)|May 29, 2024

Advanced Surface Polishing TechniquesEngineering3 citations

TL;DR

This paper introduces a differentiable, thermodynamically consistent framework for learning high-dimensional free energy models in molecular coarse graining by embedding explicit temperature and parameter dependence into machine learning models. By leveraging exact thermodynamic relationships—particularly the connection between free energy, ensemble-averaged forces, and potential energy—the method improves learning efficiency and accuracy beyond standard force-matching, enabling better sampling statistics and exact Boltzmann distribution recovery for complex, multimodal systems.

ABSTRACT

We present a differentiable formalism for learning free energies that is capable of capturing arbitrarily complex model dependencies on coarse-grained coordinates and finite-temperature response to variation of general system parameters. This is done by endowing models with explicit dependence on temperature and parameters and by exploiting exact differential thermodynamic relationships between the free energy, ensemble averages, and response properties. Formally, we derive an approach for learning high-dimensional cumulant generating functions using statistical estimates of their derivatives, which are observable cumulants of the underlying random variable. The proposed formalism opens ways to resolve several outstanding challenges in bottom-up molecular coarse graining dealing with multiple minima and state dependence. This is realized by using additional differential relationships in the loss function to significantly improve the learning of free energies, while exactly preserving the Boltzmann distribution governing the corresponding fine-grain all-atom system. As an example, we go beyond the standard force-matching procedure to demonstrate how leveraging the thermodynamic relationship between free energy and values of ensemble averaged all-atom potential energy improves the learning efficiency and accuracy of the free energy model. The result is significantly better sampling statistics of structural distribution functions. The theoretical framework presented here is demonstrated via implementations in both kernel-based and neural network machine learning regression methods and opens new ways to train accurate machine learning models for studying thermodynamic and response properties of complex molecular systems.

Motivation & Objective

• To address the challenge of learning accurate, high-dimensional free energy models in molecular coarse graining that capture complex, multimodal distributions and state-dependent behavior.
• To overcome limitations of standard force-matching and iterative Boltzmann inversion by incorporating thermodynamic consistency into the learning objective.
• To enable efficient and accurate learning of potential of mean force (PMF) models by leveraging both force and all-atom potential energy data through a joint loss function.
• To ensure exact preservation of the Boltzmann distribution governing the all-atom system, maintaining thermodynamic consistency across scales.
• To demonstrate the framework’s effectiveness using both kernel-based (Gaussian process) and neural network models in realistic coarse-grained simulations.

Proposed method

• The method introduces temperature- and parameter-dependent descriptors using Chebyshev polynomial expansions to embed thermodynamic variables directly into the model architecture.
• It formulates the free energy learning as a regression of cumulant generating functions via statistical estimates of their derivatives—observable cumulants such as forces and potential energy.
• A differentiable loss function is constructed using exact differential thermodynamic relationships, including the link between free energy, ensemble-averaged forces, and the derivative of the free energy with respect to system parameters.
• The framework integrates both force and all-atom potential energy data into a joint loss function, improving model generalization and convergence.
• The approach is implemented in both Gaussian process and neural network models, with temperature embeddings in the network architecture to capture finite-temperature response.
• The method ensures exact Boltzmann consistency by preserving the thermodynamic relationship between the coarse-grained model and the underlying all-atom system.

Experimental results

Research questions

• RQ1 Can a machine learning model for coarse-grained free energy learning be made differentiable and thermodynamically consistent by embedding explicit temperature and parameter dependence?
• RQ2 How does incorporating the thermodynamic relationship between free energy and ensemble-averaged all-atom potential energy improve learning efficiency and accuracy compared to standard force-matching?
• RQ3 To what extent can this framework capture complex, multimodal free energy landscapes with multiple minima and strong state dependence?
• RQ4 Does the joint optimization of forces and potential energy in the loss function lead to better sampling statistics of structural distribution functions?
• RQ5 Can the proposed method maintain exact Boltzmann consistency with the all-atom reference system while enabling faster and more accurate training than existing methods like REM or IBI?

Key findings

• The thermodynamically informed learning framework significantly improves sampling statistics of structural distribution functions compared to standard force-matching, particularly in systems with complex, multimodal free energy landscapes.
• Incorporating both force and all-atom potential energy data into the loss function leads to faster convergence and higher accuracy than force-only training, as validated by reduced force mean squared error and improved structural distribution agreement.
• The method preserves the exact Boltzmann distribution of the all-atom system, ensuring thermodynamic consistency between coarse-grained and all-atom ensembles.
• The use of temperature-dependent descriptors based on Chebyshev polynomials enables accurate modeling of finite-temperature response and enhances generalization across different thermodynamic states.
• The framework is successfully implemented in both kernel-based (Gaussian process) and neural network models, demonstrating broad applicability across different machine learning architectures.
• The joint loss function, which includes both force and energy terms, acts as a regularizer that improves model robustness and generalization, especially in data-scarce regions of the free energy surface.