Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework without Data

TL;DR

The paper addresses the bottleneck of data-dependent coarse-graining by proposing a fully data-free, energy-based framework that targets the all-atom Boltzmann distribution using only interatomic potentials. It decouples coordinates into slow, multimodal latent variables $\mathbf{z}$ and fast, unimodal $\mathbf{X}$ through a bijective map $\boldsymbol{f}_{\boldsymbol{\phi}}$, trained by minimizing the reverse KL divergence between a normalizing-flow–based $q_{\boldsymbol{\theta}}(\mathbf{X},\mathbf{z})$ and the transformed Boltzmann density $p_{\boldsymbol{\phi}}(\mathbf{X},\mathbf{z};\beta)$. An adaptive tempering scheme across $\beta$ mitigates mode-seeking behavior and ensures coverage of metastable states, enabling one-shot generation of fully atomistic configurations and back-mapping from coarse to fine coordinates. The approach is validated on a 2D double-well, a Gaussian mixture model, and alanine dipeptide, demonstrating accurate multimodal sampling, physically meaningful coarse representations, and faithful reproduction of structural observables without relying on precomputed MD trajectories. The method offers a principled, scalable alternative to traditional CG techniques with potential extensions to nonlinear mappings and equivariant, physics-informed flows for broader molecular systems.

Abstract

Coarse-grained (CG) models provide an effective route to reducing the complexity of molecular simulations (MD), but conventional approaches depend heavily on long all-atom MD trajectories to adequately sample configurational space. This data dependence limits accuracy and generalizability, as unvisited configurations remain excluded from the resulting CG models. We introduce a fully data-free, generative framework for CG that directly targets the all-atom Boltzmann distribution. The model defines a structured latent space comprising slow collective variables, associated with multimodal marginal densities capturing metastable states, and fast variables, represented through simple, unimodal conditional distributions. A learnable, bijective map from latent space to atomistic coordinates enables the automatic and accurate reconstruction of molecular structures. Training relies solely on the interatomic potential and minimizes the reverse Kullback-Leibler (KL) divergence via an energy-based objective. To stabilize optimization and ensure mode coverage, we employ an adaptive tempering scheme that promotes the exploration of diverse configurations. Once trained, the model can generate independent, one-shot equilibrium samples at full atomic resolution. Validation on two synthetic systems, a double-well potential and a Gaussian mixture model, as well as on the benchmark alanine dipeptide, demonstrates that the method captures all relevant modes of the Boltzmann distribution, reconstructs atomic configurations, and automatically learns physically meaningful CG representations. These results suggest a promising, data-free alternative to traditional CG techniques, offering both a principled approach to addressing the long-standing "chicken-and-egg" challenge in coarse-graining and an effective solution to the back-mapping problem by enabling accurate reconstruction of all-atom configurations.

Figures (17)

• Figure 1: Schematic illustration of the proposed generative framework. Two sets of latent coordinates are identified: a) $\mathbf{z}$ with a multimodal, learnable density $q_{\boldsymbol{\theta}}(\mathbf{z})$ corresponding to "slow" DOFs, and b) $\mathbf{X}$ with a unimodal, learnable, conditional density $q_{\boldsymbol{\theta}}(\mathbf{X}|\mathbf{z})$ corresponding to "fast" DOFs modulated by $\mathbf{z}$. These are combined in order to reconstruct the all-atom DOFs $\mathbf{x}$ through the learnable diffeomorphism $\boldsymbol{f}_{\boldsymbol{\phi}}$.
• Figure 2: Contour lines of the reverse KL-divergence in Equation \ref{['eq:klflow']} for different value pairs $(a_1,a_2)$ of the right stochastic matrix $\boldsymbol{A}_{\boldsymbol{\phi}}$ in Equation \ref{['eq:DWtransf']}. The red $\times$ represents the starting values, and the red line represents the values during training. The final learned transformation, based on $\boldsymbol{A}_{\boldsymbol{\phi}}^{-1}$, is $z = 1.03 \cdot x_1 - 0.04 \cdot x_2$ and $X = -0.03 \cdot x_1 + 1.03 \cdot x_2$.
• Figure 3: Left: Effective potential (PMF) $\beta U(x_1, x_2=0)=- \log p(x_1,x_2=0)$ (orange) and the predicted $U_{\boldsymbol{\theta}}(x_1,x_2=0)=-\log q_{\boldsymbol{\theta}}(\boldsymbol{f}_{\boldsymbol{\phi}}^{-1}(x_1,x_2=0)) + \log K_{\boldsymbol{\phi}}$ (blue) during training. Right: Histogram of samples from the marginal $p(x_1)$ (orange) and the predicted model $q_{\boldsymbol{\theta}}(X,z)$ (blue). Results are shown at the inverse temperatures $(\mathbf{a}) ~\beta\approx0$, $(\mathbf{b}) ~\beta=0.2$, $(\mathbf{c}) ~\beta=0.6$, and $(\mathbf{d}) ~\beta=1$.
• Figure 4: (Left) Histogram of all-atom samples from the target Boltzmann of Equation \ref{['eq:DoubleWell']} and (Right) from the energy-trained approximation $q_{\boldsymbol{\theta}}(\mathbf{X},\mathbf{z})$ ($\beta_{\mathrm{target}}=1$).
• Figure 5: $(\mathbf{a})$ Right stochastic matrix $\boldsymbol{A}_{\boldsymbol{\phi}}$ at (left) initialization and (right) at target inverse temperature $\beta=1$. $(\mathbf{b})$ Inverse matrix $\boldsymbol{A}_{\boldsymbol{\phi}}^{-1}$ at $\beta=1$.