Sampling Boltzmann distributions via normalizing flow approximation of transport maps

TL;DR

This paper establishes a rigorous mathematical foundation for using normalizing flows to approximate and sample from high-dimensional Boltzmann distributions in molecular dynamics, demonstrating theoretical guarantees and practical effectiveness.

Contribution

It proves the existence of a normalizing flow approximation for Boltzmann distributions with low regularity and validates the approach through numerical simulations.

Findings

Normalizing flows can approximate Boltzmann distributions with small Wasserstein error.

The approach captures both equilibrium distributions and metastable dynamics.

Numerical results confirm the theoretical predictions.

Abstract

In a celebrated paper \cite{noe2019boltzmann}, No\'e, Olsson, K\"ohler and Wu introduced an efficient method for sampling high-dimensional Boltzmann distributions arising in molecular dynamics via normalizing flow approximation of transport maps. Here, we place this approach on a firm mathematical foundation. We prove the existence of a normalizing flow between the reference measure and the true Boltzmann distribution up to an arbitrarily small error in the Wasserstein distance. This result covers general Boltzmann distributions from molecular dynamics, which have low regularity due to the presence of interatomic Coulomb and Lennard-Jones interactions. The proof is based on a rigorous construction of the Moser transport map for low-regularity endpoint densities and approximation theorems for neural networks in Sobolev spaces. Numerical simulations for a simple model system and for the alanine dipeptide molecule confirm that the true and generated distributions are close in the Wasserstein distance. Moreover we observe that the RealNVP architecture does not just successfully capture the equilibrium Boltzmann distribution but also the metastable dynamics.