2025/06/22 by Jaemoo Choi, Choi, Jaemoo, Yongxin Chen +5 · 7 citations
Computer Science · #Artificial Intelligence (cs.AI) #Bayesian Methods and Mixture Models #FOS: Computer and information sciences #Machine Learning (cs.LG)
paper · pdf · doi:10.48550/arxiv.2506.18165
openalex publication_date 2025/06/22 · openalex created_date 2025/10/09 · openalex updated_date 2026/07/28
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.