Improving sampling by modifying the effective diffusion

TL;DR

This paper introduces a method to optimize diffusion matrices in Langevin-based samplers using collective variables, enhancing exploration efficiency in high-dimensional Bayesian and physical systems.

Contribution

It proposes a scalable approach to optimize diffusion matrices via collective variables, improving sampling efficiency without costly high-dimensional optimization.

Findings

Enhanced sampling efficiency demonstrated in metastable state transitions

Implementation of optimized diffusion in Langevin and Hamiltonian Monte Carlo algorithms

Numerical results show faster convergence and better exploration in high-dimensional systems

Abstract

Markov chain Monte Carlo samplers based on discretizations of (overdamped) Langevin dynamics are commonly used in the Bayesian inference and computational statistical physics literature to estimate high-dimensional integrals. One can introduce a non-constant diffusion matrix to precondition these dynamics, and recent works have optimized it in order to improve the rate of convergence to stationarity by overcoming entropic and energy barriers. However, the introduced methodologies to compute these optimal diffusions are generally not suited to high-dimensional settings, as they rely on costly optimization procedures. In this work, we propose to optimize over a class of diffusion matrices, based on one-dimensional collective variables (CVs), to help the dynamics explore the latent space defined by the CV. The form of the diffusion matrix is chosen in order to obtain an efficient effective diffusion in the latent space. We describe how this class of diffusion matrices can be constructed and learned during the simulation. We provide implementations of the Metropolis--Adjusted Langevin Algorithm and Riemann Manifold (Generalized) Hamiltonian Monte Carlo algorithms, and discuss numerical optimizations in the case when the CV depends only on a few degrees of freedom of the system. We illustrate the efficiency gains by computing mean transition durations between two metastable states of a dimer in a solvent.