AquaGen: Scaling generative models to molecular dynamics precision on thousands of atoms

2026-07-03Machine Learning

Machine Learning
AI summary

The authors introduce AquaGen, a new computational tool that creates detailed molecular setups including water molecules much faster than traditional simulation methods. Unlike older models that simplify or skip parts of the system, AquaGen works with full atomic detail and actual solvent, allowing for accurate energy calculations and predictions. They showed that AquaGen can estimate how molecules interact with water just as accurately as standard methods but several times faster. This makes it easier to study important chemical properties and could help improve drug and material design in the future.

Generative modelBoltzmann distributionMolecular dynamics (MD)Explicit solventPeriodic boundary conditionsHydration free energyForce fieldAlchemical simulationsEnsemble averagesDrug discovery
Authors
Emmanuel Bengio, Sanjeev Raja, Yui Tik Pang, Kerstin Klaeser, Cristian Gabellini, Nikhil Shenoy, Francesco Di Giovanni, Prudencio Tossou
Abstract
We present AquaGen, the first all-atom, explicit solvent, periodic-boundary-condition-aware generative model that produces molecular configurations from the Boltzmann distribution at a fraction of the cost of molecular dynamics (MD). This is in contrast with existing generative models that remove degrees of freedom by operating on coarse-grained, vacuum, or implicit solvent systems. Operating at this resolution allows for post-processing through force field energy evaluations and MD simulations, and enables the prediction of relevant properties in a gray-box manner (as ensemble averages of potential energy evaluations over generated samples). We demonstrate the utility of this paradigm on absolute hydration free energy (AHFE), producing estimates 4-10x faster and with comparable accuracy to standard GPU-based MD. By generating uncorrelated samples from alchemical Boltzmann distributions, we create more accurate, interpretable, and refinable ensemble predictions with calibrated uncertainty estimates, unlike regression methods which are entirely black-box predictors. Our approach also yields predictable benefits from increasing train- and test-time compute, realized by scaling model size and generating more samples, respectively. We believe that this approach demonstrates the utility of high-resolution ensemble generation for free energy estimation, with future potential to replace MD in tasks such as the prediction of lipophilicity, membrane permeability, or absolute binding free energy (ABFE) -- whose grounding and interpretability may be critical for the development of new drugs and materials.