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

arXiv: https://arxiv.org/abs/2607.03513 · HTML: https://mlsanity.com/paper/2607.03513/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 52 references that are in the map.
- First posted on arXiv 3 Jul 2026.

## Builds on (most cited first, 17 of 17)

- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [Guiding a Diffusion Model with a Bad Version of Itself](https://mlsanity.com/paper/2406.02507/) (2024)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [A General Framework for Inference-time Scaling and Steering of Diffusion Models](https://mlsanity.com/paper/2501.06848/) (2025)
- [Uncertainty Quantification and Deep Ensembles](https://mlsanity.com/paper/2007.08792/) (2020)
- [ChemBERTa-2: Towards Chemical Foundation Models](https://mlsanity.com/paper/2209.01712/) (2022)
- [AlphaFold Meets Flow Matching for Generating Protein Ensembles](https://mlsanity.com/paper/2402.04845/) (2024)
- [Generative Modeling of Molecular Dynamics Trajectories](https://mlsanity.com/paper/2409.17808/) (2024)
- [ET-Flow: Equivariant Flow-Matching for Molecular Conformer Generation](https://mlsanity.com/paper/2410.22388/) (2024)
- [FEAT: Free energy Estimators with Adaptive Transport](https://mlsanity.com/paper/2504.11516/) (2025)
- [Neural Thermodynamic Integration: Free Energies from Energy-based Diffusion Models](https://mlsanity.com/paper/2406.02313/) (2024)
- [Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules](https://mlsanity.com/paper/2401.04246/) (2024)
- [Solvation Free Energies from Neural Thermodynamic Integration](https://mlsanity.com/paper/2410.15815/) (2024)
- [MarS-FM: Generative Modeling of Molecular Dynamics via Markov State Models](https://mlsanity.com/paper/2509.24779/) (2025)
- [Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics](https://mlsanity.com/paper/2603.17633/) (2026)

## Abstract (arXiv)

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.
