Improving sampling by modifying the effective diffusion

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.