Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sampling Boltzmann distributions via normalizing flow approximation of transport maps

Zia Ur Rehman, Gero Friesecke

Abstract

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.

Sampling Boltzmann distributions via normalizing flow approximation of transport maps

Abstract

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.
Paper Structure (14 sections, 5 theorems, 39 equations, 2 figures)

This paper contains 14 sections, 5 theorems, 39 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Molecular dynamics at thermal equilibrium
  3. Transport maps and their lack of regularity
  4. Normalizing flow approximation of transport maps
  5. Composition structure of the Brenier and Moser maps
  6. Theoretical results
  7. Regularization
  8. Moser transport map: existence and approximation by RealNVP flow
  9. Convergence of push-forward density
  10. Main result
  11. Numerical Experiments
  12. 2D Double-well potential
  13. Alanine Dipeptide Molecular System
  14. Discussion and Conclusions

Key Result

Proposition 1

Let $\Omega\subset{\mathbb R}^{3N}$ be any open bounded set, let $\rho_0$ be any reference probability density on $\Omega$ which is continuous and strictly positive, and let $\rho_1$ be a Boltzmann distribution boltz on an open bounded set $D\subset{\mathbb R}^{3N}$ with $U$ satisfying potquali. The

Figures (2)

  • Figure 1: Comparison of real Langevin dynamics and RealNVP generated samples for the 2D double-well potential.
  • Figure 2: Comparison of true and generated $\Phi$--$\Psi$ distributions for alanine dipeptide

Theorems & Definitions (10)

  • Proposition 1: Nonexistence of Lipschitz transport maps
  • proof : Proof
  • Theorem 1: Convergence of regularized Boltzmann distributions
  • proof : Proof
  • Theorem 2: Existence and RealNVP approximation of the Moser transport map
  • proof : Proof
  • Theorem 3: Convergence of push-forward densities
  • proof : Proof
  • Theorem 4: Theoretical justification of the Boltzmann generator method
  • proof : Proof