Non-equilibrium Annealed Adjoint Sampler

TL;DR

The paper introduces NAAS, a novel diffusion sampling framework using annealed reference dynamics and stochastic optimal control, improving efficiency and stability in sampling from complex distributions.

Contribution

NAAS employs annealed non-stationary reference dynamics within a SOC framework, offering flexible algorithm design and scalable training for diffusion sampling.

Findings

Effective in sampling from energy landscapes

Improves stability and efficiency of diffusion samplers

Demonstrates success on molecular Boltzmann distributions

Abstract

Recently, there has been significant progress in learning-based diffusion samplers, which aim to sample from a given unnormalized density. Many of these approaches formulate the sampling task as a stochastic optimal control (SOC) problem using a canonical uninformative reference process, which limits their ability to efficiently guide trajectories toward the target distribution. In this work, we propose the Non-Equilibrium Annealed Adjoint Sampler (NAAS), a novel SOC-based diffusion framework that employs annealed reference dynamics as a non-stationary base SDE. This annealing structure provides a natural progression toward the target distribution and generates informative reference trajectories, thereby enhancing the stability and efficiency of learning the control. Owing to our SOC formulation, our framework can incorporate a variety of SOC solvers, thereby offering high flexibility in algorithmic design. As one instantiation, we employ a lean adjoint system inspired by adjoint matching, enabling efficient and scalable training. We demonstrate the effectiveness of NAAS across a range of tasks, including sampling from classical energy landscapes and molecular Boltzmann distributions.