Sampling Boltzmann distributions via normalizing flow approximation of transport maps
This work addresses the challenge of sampling complex molecular dynamics distributions, offering a rigorous theoretical basis for an existing method, which is incremental but important for reliability in computational chemistry.
The paper provides a mathematical foundation for sampling high-dimensional Boltzmann distributions in molecular dynamics using normalizing flows, proving existence of such flows with arbitrarily small error in Wasserstein distance and confirming closeness in numerical simulations for model systems and alanine dipeptide.
In a celebrated paper \cite{noe2019boltzmann}, Noé, Olsson, Köhler 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.