Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations

TL;DR

This work addresses the challenge of efficiently sampling the Boltzmann distribution for molecular systems by introducing a conditional normalizing flow that maps latent variables to fine-grained configurations conditioned on a coarse-grained (CG) space, coupled with an active-learning loop in CG space to refine the PMF. Training by energy allows simulation-free learning of the conditional distribution, while an ensemble of PMF models guides targeted exploration of high-uncertainty CG configurations. The authors demonstrate substantial data and time savings on Müller-Brown and alanine dipeptide benchmarks, achieving PMF accuracy comparable to or better than MD with orders-of-magnitude fewer energy evaluations and speedups up to $216.2\times$ over MD. This approach offers a scalable framework for ML-assisted coarse-graining and unnormalized-density sampling, with potential extensions to larger systems, nonlinear CG mappings, and equivariant normalizing-flow architectures.

Abstract

Efficient sampling of the Boltzmann distribution of molecular systems is a long-standing challenge. Recently, instead of generating long molecular dynamics simulations, generative machine learning methods such as normalizing flows have been used to learn the Boltzmann distribution directly, without samples. However, this approach is susceptible to mode collapse and thus often does not explore the full configurational space. In this work, we address this challenge by separating the problem into two levels, the fine-grained and coarse-grained degrees of freedom. A normalizing flow conditioned on the coarse-grained space yields a probabilistic connection between the two levels. To explore the configurational space, we employ coarse-grained simulations with active learning which allows us to update the flow and make all-atom potential energy evaluations only when necessary. Using alanine dipeptide as an example, we show that our methods obtain a speedup to molecular dynamics simulations of approximately 15.9 to 216.2 compared to the speedup of 4.5 of the current state-of-the-art machine learning approach.

Figures (11)

• Figure 1: (a) The conditional normalizing flow transforms the latent variable $z$ of the latent distribution to the target configuration $x_\text{int}$ conditioned on the CG configuration $s$. (b) Illustration of the iterative three-step active learning cycle.
• Figure 2: Visualization of two exemplary iterations of the AL workflow applied to the Müller-Brown system. Bottom: PMF and its standard deviation. Training points from previous AL iterations are marked as black "x" at the top of the PMF, and new high-error points added in the current AL iteration are marked as red "x". Top: Backmapped potential energy $\ln q_{S_\perp}(s_\perp \mid s; \theta) p^\text{CG}(s)$ from cascaded sampling of the PMF and sampling of the flow. The blue axis represents the 1D CG coordinate $s$. The fine-grained coordinate $s_\perp$ is orthogonal to the CG coordinate. See Figure \ref{['fig:SI:mb_steps']} in SI for a visualization of all iterations of this experiment.
• Figure 3: (a) Ground truth PMF of dihedral angles $\phi$ and $\psi$ of alanine dipeptide from MD test dataset with 2.3e10 steps. (b) PMF from grid conditioning experiment after 2.4e7 steps. (c-e) Three exemplary AL steps for the alanine dipeptide system. Top: PMF of $\phi$ and $\psi$ at the end of the AL iteration. Bottom: Standard deviation of the PMF. The contour line of the threshold in the standard deviation when sampling new points ($0.2 \, k_\mathrm{B} T$) is drawn using a black line. Newly added points after sampling using the PMF of this iteration are shown as red dots.
• Figure 4: Forward KLD of the PMF of alanine dipeptide as a function of the number of potential energy evaluations. For grid conditioning and Flow AIS Bootstrap we additionally show the standard error.
• Figure 5: Müller-Brown potential. The blue axis represents the 1D CG coordinate $s$. The fine-grained coordinate $s_\perp$ is orthogonal to the CG coordinate.