Learning Boltzmann Generators via Constrained Mass Transport

TL;DR

This work tackles the challenge of sampling from high-dimensional, unnormalized targets by improving Boltzmann generators through Constrained Mass Transport (CMT). CMT formulates a sequence of intermediate densities connecting a tractable base to the target while enforcing both KL-distance and entropy-decay constraints, mitigating mass teleportation and mode collapse. The method is instantiated with normalizing flows, using a forward KL objective and a dual optimization scheme to learn intermediate densities efficiently, reusing samples via replay buffers. Empirically, CMT outperforms state-of-the-art variational BG methods on multiple alanine-based benchmarks and introduces the ELIL tetrapeptide, achieving over 2.5x higher effective sample size and better mode coverage without relying on MD samples. The results underscore the practical impact of constrained annealing paths for robust, scalable Boltzmann sampling in complex molecular systems.

Abstract

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.

Figures (20)

• Figure 1: Illustration of the annealing paths (AP) obtained by solving the variational problems \ref{['eq: tr problem']}, \ref{['eq: entropy problem']}, or \ref{['eq: tr entropy problem']}. Trust-region–based optimization \ref{['eq: tr problem']} mitigates the irregularities of naive schedules (e.g., the linear schedule), but the resulting geometric AP suffers from mass teleportation as the right mode of the target distribution $p$ emerges without overlap with earlier intermediate densities. Constraining the entropy decay between successive densities \ref{['eq: entropy problem']} prevents mass teleportation, yet fails to guarantee sufficient overlap between the initial distribution $q_0$ and subsequent intermediate densities. In contrast, combining both constraints \ref{['eq: tr entropy problem']} yields APs that both maintain overlap between successive densities and avoid mass teleportation.
• Figure 2: Results for all systems of varying dimensionality $d$. Evaluation criteria include the number of target evaluations (Target evals), the evidence upper bound (EUBO), the reverse effective sample size (ESS) and the average total variation distance to the Ramachandran plots generated from molecular dynamics samples (Ram TV). Details on all metrics can be found in \ref{['appendix:metrics']}. Each value is shown as the mean $\pm$ standard error over four independent runs. The best results are highlighted in bold, except for the forward KL and reverse KL. Reverse KL is prone to mode collapse, which makes ESS values not directly comparable, and forward KL is trained from samples rather than from energy.
• Figure 3: Impact of the trust-region and entropy constraint visualized on alanine hexapeptide. \ref{['fig:hexa_diff_constraints_entropy_progression']} visualizes the model entropy over the course of the training. Analogously, \ref{['fig:hexa_diff_constraints_iw_ess']} shows the importance-weight between successive intermediate densities. \ref{['fig:hexa_diff_constraints_final_log_likelihood', 'fig:hexa_diff_constraints_final_ess']} depict the final log-likelihood and to the target density, respectively. The variants in \ref{['fig:hexa_diff_constraints_final_ess']} marked with "$\textcolor{red}{\bigstar}$" exhibit visible mode-collapse on the Ramachandran plots; see \ref{['fig:diff_constraints_ramachandrans_hexa']}. The is therefore not directly comparable to the other methods. We denote by $\hat{q}_i$ the variational approximation of the intermediate density $q_i$.
• Figure 4: Ramachandran plots for alanine hexapeptide with trust-region and entropy constraints selectively enabled or disabled. Using a single or no constraint leads to mode collapse, whereas combining both constraints avoids it. See \ref{['appendix:metrics']} for details.
• Figure 5: Comparison of Ramachandran plots of backbone dihedral angle pairs obtained with different methods. See \ref{['appendix:metrics']} for details on Ramachandran plots.

Theorems & Definitions (7)

• Proposition 2.1: Optimal intermediate trust-region densities
• Proposition 2.2: Optimal intermediate densities for entropy constraint
• Proposition 2.3: Optimal intermediate densities for entropy and trust-region constraint
• Theorem 2.4: Annealing paths
• proof : Proof of \ref{['prop: optimal tr densities', 'prop: optimal ent densities', 'prop: optimal tr-ent densities']}
• proof : Proof of \ref{['prop: annealing paths']}
• proof : Proof of uniqueness and tightness of the trust-region solution