Skipping the Replica Exchange Ladder with Normalizing Flows

TL;DR

This paper introduces learned replica exchange (LREX), a novel method combining replica exchange with normalizing flows to efficiently sample molecular systems with rare events, reducing computational costs and improving over existing approaches.

Contribution

The paper presents LREX, a new approach that skips intermediate replicas by training normalizing flows for direct configuration exchange, enhancing efficiency in molecular sampling.

Findings

LREX reduces computational cost compared to standard replica exchange.

LREX outperforms Boltzmann generators in sampling efficiency.

Application to molecular systems demonstrates significant improvements.

Abstract

We combine replica exchange (parallel tempering) with normalizing flows, a class of deep generative models. These two sampling strategies complement each other, resulting in an efficient strategy for sampling molecular systems characterized by rare events, which we call learned replica exchange (LREX). In LREX, a normalizing flow is trained to map the configurations of the fastest-mixing replica into configurations belonging to the target distribution, allowing direct exchanges between the two without the need to simulate intermediate replicas. This can significantly reduce the computational cost compared to standard replica exchange. The proposed method also offers several advantages with respect to Boltzmann generators that directly use normalizing flows to sample the target distribution. We apply LREX to some prototypical molecular dynamics systems, highlighting the improvements over previous methods.