Non-equilibrium Annealed Adjoint Sampler
This addresses a bottleneck in sampling methods for researchers in machine learning and computational science, offering a more scalable alternative to existing approaches.
The paper tackles the problem of high variance and limited scalability in diffusion samplers by introducing the Non-equilibrium Annealed Adjoint Sampler (NAAS), which uses annealed reference dynamics without importance sampling, achieving efficient and scalable training across tasks like sampling from energy landscapes and molecular Boltzmann distributions.
Recently, there has been significant progress in learning-based diffusion samplers, which aim to sample from a given unnormalized density. These methods typically follow one of two paradigms: (i) formulating sampling as an unbiased stochastic optimal control (SOC) problem using a canonical reference process, or (ii) refining annealed path measures through importance-weighted sampling. Although annealing approaches have advantages in guiding samples toward high-density regions, reliance on importance sampling leads to high variance and limited scalability in practice. In this paper, we introduce the \textbf{Non-equilibrium Annealed Adjoint Sampler (NAAS)}, a novel SOC-based diffusion sampler that leverages annealed reference dynamics without resorting to importance sampling. NAAS employs a lean adjoint system inspired by adjoint matching, enabling efficient and scalable training. We demonstrate the effectiveness of our approach across a range of tasks, including sampling from classical energy landscapes and molecular Boltzmann distribution.