# Learning Boltzmann Generators via Constrained Mass Transport

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 69 of its 114 references that are in the map.
- First posted on arXiv 21 Oct 2025.

## Builds on (most cited first, 30 of 69)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Constrained Policy Optimization](https://mlsanity.com/paper/1705.10528/) (2017)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Torsional Diffusion for Molecular Conformer Generation](https://mlsanity.com/paper/2206.01729/) (2022)
- [Projection-Based Constrained Policy Optimization](https://mlsanity.com/paper/2010.03152/) (2020)
- [Understanding the impact of entropy on policy optimization](https://mlsanity.com/paper/1811.11214/) (2018)
- [Equivariant flow matching](https://mlsanity.com/paper/2306.15030/) (2023)
- [Flow Annealed Importance Sampling Bootstrap](https://mlsanity.com/paper/2208.01893/) (2022)
- [Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation](https://mlsanity.com/paper/1708.05144/) (2017)
- [Normalizing Flows on Tori and Spheres](https://mlsanity.com/paper/2002.02428/) (2020)
- [AlphaFold Meets Flow Matching for Generating Protein Ensembles](https://mlsanity.com/paper/2402.04845/) (2024)
- [Normalizing Flows are Capable Generative Models](https://mlsanity.com/paper/2412.06329/) (2024)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [Particle Denoising Diffusion Sampler](https://mlsanity.com/paper/2402.06320/) (2024)
- [DIME:Diffusion-Based Maximum Entropy Reinforcement Learning](https://mlsanity.com/paper/2502.02316/) (2025)

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

- [Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization](https://mlsanity.com/paper/2602.03729/) (2026)
- [Learning Structure, Energy, and Dynamics: A Survey of Artificial Intelligence for Protein Dynamics](https://mlsanity.com/paper/2604.25244/) (2026)
- [Autoregressive Boltzmann Generators](https://mlsanity.com/paper/2606.27361/) (2026)
- [Coarse-Grained Boltzmann Generators](https://mlsanity.com/paper/2602.10637/) (2026)
- [Scalable Inference-Time Annealing with Surrogate Likelihood Estimators](https://mlsanity.com/paper/2605.31498/) (2026)
- [Path Integral Value Matching for Linear Quadratic Stochastic Optimal Control](https://mlsanity.com/paper/2608.10777/) (2026)
- [Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation](https://mlsanity.com/paper/2609.09473/) (2026)
- [FERPO: Forward Entropy-Regularized Policy Optimization](https://mlsanity.com/paper/2610.02198/) (2026)
- [Looking Under the Streetlight: Evaluation in Generative Molecular Dynamics](https://mlsanity.com/paper/2610.12518/) (2026)

## Abstract (arXiv)

Efficient sampling from high-dimensional and multimodal unnormalized probability distributions is a central challenge in many areas of science and machine learning. We focus on Boltzmann generators (BGs) that aim to sample the Boltzmann distribution of physical systems, such as molecules, at a given temperature. Classical variational approaches that minimize the reverse Kullback-Leibler divergence are prone to mode collapse, while annealing-based methods, commonly using geometric schedules, can suffer from mass teleportation and rely heavily on schedule tuning. We introduce Constrained Mass Transport (CMT), a variational framework that generates intermediate distributions under constraints on both the KL divergence and the entropy decay between successive steps. These constraints enhance distributional overlap, mitigate mass teleportation, and counteract premature convergence. Across standard BG benchmarks and the here introduced ELIL tetrapeptide, the largest system studied to date without access to samples from molecular dynamics, CMT consistently surpasses state-of-the-art variational methods, achieving more than 2.5x higher effective sample size while avoiding mode collapse.