# Scalable Normalizing Flows Enable Boltzmann Generators for Macromolecules

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 46 references that are in the map.
- First posted on arXiv 8 Jan 2024.

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

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Transformers without Tears: Improving the Normalization of Self-Attention](https://mlsanity.com/paper/1910.05895/) (2019)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Normalizing Flows on Tori and Spheres](https://mlsanity.com/paper/2002.02428/) (2020)
- [Protein structure generation via folding diffusion](https://mlsanity.com/paper/2209.15611/) (2022)
- [Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics](https://mlsanity.com/paper/2302.01170/) (2023)
- [Smooth Normalizing Flows](https://mlsanity.com/paper/2110.00351/) (2021)
- [Densely connected normalizing flows](https://mlsanity.com/paper/2106.04627/) (2021)
- [Normalizing flows for atomic solids](https://mlsanity.com/paper/2111.08696/) (2021)
- [Flow-matching -- efficient coarse-graining of molecular dynamics without forces](https://mlsanity.com/paper/2203.11167/) (2022)
- [Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-Scale Graph Networks](https://mlsanity.com/paper/2204.10348/) (2022)

## Built on by (most cited first, 10 of 10)

- [Transferable Boltzmann Generators](https://mlsanity.com/paper/2406.14426/) (2024)
- [Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models](https://mlsanity.com/paper/2506.17139/) (2025)
- [BoltzNCE: Learning Likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation](https://mlsanity.com/paper/2507.00846/) (2025)
- [JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensembles](https://mlsanity.com/paper/2410.14621/) (2024)
- [Path Gradients after Flow Matching](https://mlsanity.com/paper/2505.10139/) (2025)
- [Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization](https://mlsanity.com/paper/2602.03729/) (2026)
- [Verlet Flows: Exact-Likelihood Integrators for Flow-Based Generative Models](https://mlsanity.com/paper/2405.02805/) (2024)
- [Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework without Data](https://mlsanity.com/paper/2504.20940/) (2025)
- [Learning Structure, Energy, and Dynamics: A Survey of Artificial Intelligence for Protein Dynamics](https://mlsanity.com/paper/2604.25244/) (2026)
- [AquaGen: Scaling generative models to molecular dynamics precision on thousands of atoms](https://mlsanity.com/paper/2607.03513/) (2026)

## Abstract (arXiv)

The Boltzmann distribution of a protein provides a roadmap to all of its functional states. Normalizing flows are a promising tool for modeling this distribution, but current methods are intractable for typical pharmacological targets; they become computationally intractable due to the size of the system, heterogeneity of intra-molecular potential energy, and long-range interactions. To remedy these issues, we present a novel flow architecture that utilizes split channels and gated attention to efficiently learn the conformational distribution of proteins defined by internal coordinates. We show that by utilizing a 2-Wasserstein loss, one can smooth the transition from maximum likelihood training to energy-based training, enabling the training of Boltzmann Generators for macromolecules. We evaluate our model and training strategy on villin headpiece HP35(nle-nle), a 35-residue subdomain, and protein G, a 56-residue protein. We demonstrate that standard architectures and training strategies, such as maximum likelihood alone, fail while our novel architecture and multi-stage training strategy are able to model the conformational distributions of protein G and HP35.