Improving molecular sampling with stepwise corrected flow models
Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
Machine Learning
Summary
Simulating large molecular systems accurately and efficiently is very challenging. The authors introduce a new method called NF²M that gradually improves predicted molecular structures at each step, rather than only fixing errors at the end like most methods. This approach uses a special model that can exactly calculate how likely proposed structures are, allowing better corrections during generation. When tested on molecules like peptides, this method generated better and more efficient samples. This helps scientists better understand molecular behavior by creating realistic molecular examples more quickly.
What this means in practice
- •For computational chemists: Generate higher-quality equilibrium samples of molecular systems with improved efficiency for simulation studies.
- •For pharmaceutical modelers: Use enhanced sampling of molecular configurations to better predict drug-target interactions in early-stage drug design.
Authors
Louis Grenioux, RuiKang OuYang, Luhuan Wu
Abstract
Generating independent, equilibrium samples of molecular systems at scale remains a central obstacle in computational statistical mechanics. Boltzmann Generators address this by pairing a generative model with importance sampling to obtain consistent samples from the target distribution. We introduce Normalizing Flow Flow Maps (NF$^2$M), which combines the strengths of recent stochastic flow maps with the tractability of classic normalizing flows to build a Boltzmann Generator. Unlike most methods, which correct the generative model only at the end, NF$^2$M reweighs each denoising transition as generation proceeds, avoiding wasted compute on trajectories that are ultimately discarded. At each denoising step, a conditional normalizing flow proposes clean configurations given the current noisy state (a simpler task than sampling directly from the target) and its exact likelihood enables correcting each proposal toward the true denoising transition of the target Boltzmann distribution. This is in contrast to most existing methods, whose likelihoods are approximate or expensive to evaluate, undermining the statistical reliability of the correction. We establish consistency of the corrected transitions and bound how approximation errors propagate through the sampling chain. We evaluate NF$^2$M on peptide systems, demonstrating improved sampling efficiency and sample quality.