Non-equilibrium Annealed Adjoint Sampler

2025-06-22

AI Summary

Existing learning-based diffusion sampling methods for unnormalized densities either rely on unbiased stochastic optimal control (SOC) frameworks or employ importance-weighted annealing paths—though the latter facilitates guidance toward high-density regions, it suffers from high variance and scalability limitations. This paper proposes a novel SOC framework that introduces a nonequilibrium annealing adjoint system, integrating annealing reference dynamics with adjoint-matching principles to enable efficient density-guided sampling without importance sampling. The lightweight adjoint system substantially reduces estimator variance, enhancing training stability and scalability. Experiments on classical energy landscape sampling and molecular Boltzmann distribution modeling demonstrate that our method achieves superior sampling efficiency and generalization performance compared to state-of-the-art approaches.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Machine Learning: Learning with Manifolds

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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 extbf{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.

Problem

• Develops a scalable diffusion sampler without importance sampling
• Addresses high variance in annealed path measure methods
• Improves sampling from unnormalized densities like energy landscapes

Innovation

• Novel SOC-based diffusion sampler without importance sampling
• Uses annealed reference dynamics for better sampling
• Employs lean adjoint system for efficient training