Non-equilibrium Annealed Adjoint Sampler
Jaemoo Choi, Yongxin Chen, Molei Tao, Guan-Horng Liu
Recently, there has been significant progress in learning-based diffusion samplers, which aim to sample from 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 (IWS). 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 **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.