Learning Boltzmann Generators via Constrained Mass Transport

TL;DR

This paper introduces Constrained Mass Transport, a new variational framework for Boltzmann generators that improves sampling efficiency and diversity in high-dimensional, multimodal distributions, especially in molecular systems.

Contribution

The paper proposes Constrained Mass Transport, a novel approach that constrains KL divergence and entropy decay to enhance Boltzmann generator performance and prevent mode collapse.

Findings

CMT outperforms existing methods on standard benchmarks.

Achieves over 2.5x higher effective sample size.

Successfully applied to large molecular system without prior samples.

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.