Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sampling from Boltzmann densities with physics informed low-rank formats

Paul Hagemann, Janina Schütte, David Sommer, Martin Eigel, Gabriele Steidl

TL;DR

This paper tackles sampling from unnormalized Boltzmann densities by learning a velocity field $v_t$ that transports samples along an annealing path from a latent distribution $p_0$ to a target $p_1$ through the continuity equation $ rac{\partial p_t}{\partial t} + \nabla\cdot(p_t v_t)=0$. It introduces a functional tensor train (FTT) representation for $v_t$, optimized with alternating linear scheme (ALS) on a low-rank TT manifold and enhanced with Langevin and resampling steps to improve mode exploration and avoid teleportation near the final annealing time. An adaptive TT-rank strategy based on exponential moving averages and TT-SVD provides automatic complexity control, while the use of an $H^2$-orthonormal Fourier basis improves PDE solvability. Numerical experiments on 2D Gaussian mixtures and the Many Well Problem show that the TT flow benefits from stochastic steps, achieving accurate target sampling in higher dimensions and complex multimodal landscapes. The approach offers a scalable, physics-informed alternative to classical MCMC and flow-based sampling for high-dimensional, multimodal distributions.

Abstract

Our method proposes the efficient generation of samples from an unnormalized Boltzmann density by solving the underlying continuity equation in the low-rank tensor train (TT) format. It is based on the annealing path commonly used in MCMC literature, which is given by the linear interpolation in the space of energies. Inspired by Sequential Monte Carlo, we alternate between deterministic time steps from the TT representation of the flow field and stochastic steps, which include Langevin and resampling steps. These adjust the relative weights of the different modes of the target distribution and anneal to the correct path distribution. We showcase the efficiency of our method on multiple numerical examples.

Sampling from Boltzmann densities with physics informed low-rank formats

TL;DR

This paper tackles sampling from unnormalized Boltzmann densities by learning a velocity field that transports samples along an annealing path from a latent distribution to a target through the continuity equation . It introduces a functional tensor train (FTT) representation for , optimized with alternating linear scheme (ALS) on a low-rank TT manifold and enhanced with Langevin and resampling steps to improve mode exploration and avoid teleportation near the final annealing time. An adaptive TT-rank strategy based on exponential moving averages and TT-SVD provides automatic complexity control, while the use of an -orthonormal Fourier basis improves PDE solvability. Numerical experiments on 2D Gaussian mixtures and the Many Well Problem show that the TT flow benefits from stochastic steps, achieving accurate target sampling in higher dimensions and complex multimodal landscapes. The approach offers a scalable, physics-informed alternative to classical MCMC and flow-based sampling for high-dimensional, multimodal distributions.

Abstract

Our method proposes the efficient generation of samples from an unnormalized Boltzmann density by solving the underlying continuity equation in the low-rank tensor train (TT) format. It is based on the annealing path commonly used in MCMC literature, which is given by the linear interpolation in the space of energies. Inspired by Sequential Monte Carlo, we alternate between deterministic time steps from the TT representation of the flow field and stochastic steps, which include Langevin and resampling steps. These adjust the relative weights of the different modes of the target distribution and anneal to the correct path distribution. We showcase the efficiency of our method on multiple numerical examples.

Paper Structure

This paper contains 12 sections, 11 equations, 2 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Logarithmic Continuity Equation
  3. Low-rank tensor train decomposition
  4. Vector field learning with the alternating linear scheme
  5. Methodology
  6. Stochastic steps
  7. Adaptive TT ranks
  8. Numerical experiments
  9. Target Distributions.
  10. Methods.
  11. Conclusions
  12. Acknowledgements

Figures (2)

  • Figure 1: GM density sampled with the different methods. While the methods perform similarly on this problem, the TT flow in (b) and (c) helps to distribute the mass correctly compared to a purely stochastic approach (as can e.g. be seen for the mode on the lower left part of the domain).
  • Figure 2: Histogram corner plots showing the 1-d marginals on the diagonal and the 2-d joint distributions on the off-diagonal for the MWP $f^{4,16}_1$ for different methods (blue: competing method, orange ground truth).