Non-equilibrium Annealed Adjoint Sampler

TL;DR

This work addresses the challenge of sampling from complex unnormalized targets by casting diffusion sampling as a stochastic-optimal-control problem with non-equilibrium, annealed reference dynamics. NAAS introduces a two-stage SOC framework that first learns a tractable prior via backward optimization on an interval preceding the target, then optimizes annealed dynamics to transport that prior to the target while ensuring unbiasedness at the terminal time. The methodology leverages Adjoint Matching and Reciprocal Adjoint Matching to achieve scalable, low-variance gradient estimation and efficient training in high dimensions, with replay buffers to further enhance practicality. Empirically, NAAS delivers state-of-the-art performance on synthetic energy landscapes and molecular generation (alanine dipeptide), demonstrating improved sample fidelity, diversity, and mode coverage without relying on importance-weighted sampling during training. The work suggests a flexible, scalable framework for diffusion-based sampling that can adapt to various SOC solvers and target domains, while outlining directions for reducing computational overhead and extending to larger-scale problems.

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.

Figures (10)

• Figure 1: Visualzation of Training Process of NAAS. Given an annealing reference dynamics that provide high-quality---yet imperfect---initial samples in high-density regions of the target distribution $\nu$, NAAS learns to sample from both the target $\nu$ and the debiased source $\mu$ by alternate optimization between two control functions, $u^\theta$ and $v^\theta$, with Adjoint Matching (\ref{['eq:am-naas2']}) and Reciprocal Adjoint Matching (\ref{['eq:am-naas1']}). This iterative procedure progressively aligns the sampling path with the optimal control plan, leading to unbiased and efficient sampling from the target $\nu$.
• Figure 2: Qualitative Comparison on GMM40 (50d) and MoS (50d). For each task, we visualize the kernel density estimate (KDE) of the target distribution alongside the generated samples (shown as red dots), projected onto the first two principal axes. The samples from NAAS closely follow the structure and support of the ground-truth, demonstrating accurate mode coverage and high sample fidelity in high-dimensional settings.
• Figure 3: Visualization of training dynamics over epochs. Each runs are presented in a different color.
• Figure 3: Results for AD generation. We report KL divergence for torsions ($D_{\text{KL}}$) and Wasserstein 2-distance for energies $E(\cdot ) W_2$.
• Figure 4: Comparison of five torsions between generated and reference samples

Theorems & Definitions (11)

• Theorem 3.1: SOC-based Non-equilibrium Sampler
• Remark
• Corollary 3.2
• Lemma 3.3
• Theorem 3.4
• proof
• proof
• proof
• proof
• Proposition B.1