Sampling Boltzmann distributions via normalizing flow approximation of transport maps

Zia Ur Rehman; Gero Friesecke

cs.LGMar 15, 2026

Sampling Boltzmann distributions via normalizing flow approximation of transport maps

Authors: Zia Ur Rehman, Gero Friesecke

Organizations: Department of Mathematics, School of Computation, Information and Technology, Technical University of Munich, Boltzmannstrasse 3, 85748 Garching, Germany

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.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison of real Langevin dynamics and RealNVP generated samples for the 2D double-well potential.

Panel of Figure 2 — Figure 2: Comparison of true and generated $\Phi$ – $\Psi$ distributions for alanine dipeptide

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows
Sep 25, 2026Louis Grenioux, RuiKang OuYang, Luhuan WuMolecular Generation Normalizing Flows
Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA
Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
Oct 7, 2026Chenguang Duan, Johannes Hertrich, Gabriele SteidlWasserstein Gradient Flows
Institut für Geometrie und Praktische Mathematik, RWTH Aachen University · Institute of Computer Science, University of Göttingen · Institute of Mathematics, TU Berlin
Neural Transport Nested Sampling
Sep 24, 2026David Yallup, Will HandleyMarkov Chain Monte Carlo Molecular Dynamics Simulation
Kavli Institute for Cosmology Cambridge Institute of Astronomy, University of Cambridge, Madingley Road, Cambridge, CB3 0HA, UK

Sep 25, 2026J/K move · Enter open · S save

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Louis Grenioux, RuiKang OuYang, Luhuan Wu

Center for Computational Mathematics, Flatiron Institute, New York, NY, USA · Department of Engineering, University of Cambridge, Cambridge, UK · Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA

Molecular Generation Normalizing Flows Flow Map Learning Importance Sampling

Open paper AskPDF arXiv

Generating independent, equilibrium samples of molecular systems at scale remains a central obstacle in computational statistical mechanics. Boltzmann Generators address this by pairing a generative model with importance sampling to obtain consistent samples from the target distribution. We introduce Normalizing Flow Flow Maps (NF

^2

M), which combines the strengths of recent stochastic flow maps with the tractability of classic normalizing flows to build a Boltzmann Generator. Unlike most methods, which correct the generative model only at the end, NF

^2

M reweighs each denoising transition as generation proceeds, avoiding wasted compute on trajectories that are ultimately discarded. At each denoising step, a conditional normalizing flow proposes clean configurations given the current noisy state (a simpler task than sampling directly from the target) and its exact likelihood enables correcting each proposal toward the true denoising transition of the target Boltzmann distribution. This is in contrast to most existing methods, whose likelihoods are approximate or expensive to evaluate, undermining the statistical reliability of the correction. We establish consistency of the corrected transitions and bound how approximation errors propagate through the sampling chain. We evaluate NF

^2

M on peptide systems, demonstrating improved sampling efficiency and sample quality.

Abstract

Explore similar work

Bridging Stochastic Flow Maps and Boltzmann Generators with Normalizing Flows

Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme

Neural Transport Nested Sampling