Learning Boltzmann Generators via Constrained Mass Transport

2025-10-21

AI Summary

Sampling from high-dimensional, multimodal, unnormalized distributions—such as Boltzmann distributions in physical systems—remains challenging. Existing variational methods suffer from mode collapse, while annealing-based approaches require careful schedule tuning and incur mass transport mismatch. This paper proposes a constrained optimal transport framework: it introduces dual regularization via KL divergence and entropy decay to explicitly constrain intermediate distribution sequences, thereby suppressing mass teleportation and mode collapse; and it employs invertible neural networks to model the probability flow, enabling end-to-end, annealing-free variational inference. Evaluated on standard benchmarks and the ELIL tetrapeptide system, the method achieves over 2.5× improvement in effective sample size without requiring prior molecular dynamics samples, significantly enhancing stability, diversity, and efficiency of distribution coverage.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Security and Privacy: Large-scale security measurements

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.

Problem

• Addresses sampling from high-dimensional multimodal distributions
• Mitigates mode collapse in Boltzmann generator training
• Overcomes mass teleportation in annealing-based sampling methods

Innovation

• Constrained Mass Transport framework for Boltzmann generators
• Intermediate distributions with KL and entropy constraints
• Enhanced overlap and reduced mass teleportation in sampling