Non-equilibrium Annealed Adjoint Sampler

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Choi, Jaemoo, Chen, Yongxin, Tao, Molei, Liu, Guan-Horng
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866915635978567680
author Choi, Jaemoo
Chen, Yongxin
Tao, Molei
Liu, Guan-Horng
author_facet Choi, Jaemoo
Chen, Yongxin
Tao, Molei
Liu, Guan-Horng
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2506_18165
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Non-equilibrium Annealed Adjoint Sampler
Choi, Jaemoo
Chen, Yongxin
Tao, Molei
Liu, Guan-Horng
Machine Learning
Artificial Intelligence
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.
title Non-equilibrium Annealed Adjoint Sampler
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2506.18165