Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

2025-04-16

AI Summary

This work addresses the challenge of efficient, scalable diffusion sampling under unnormalized energy functions—such as molecular force fields or neural network potentials. We propose Adjoint Sampling, the first on-policy diffusion sampling algorithm designed specifically for unnormalized densities. Grounded in stochastic optimal control theory, our method integrates adjoint matching, energy-guided sampling, and joint optimization with neural network potentials, ensuring convergence without correction mechanisms. It is the first to support both Cartesian and dihedral coordinate representations under periodic boundary conditions and molecular symmetry constraints. Gradient updates occur at a significantly higher frequency than energy evaluations or sample generation. We demonstrate large-scale molecular conformational sampling on classical force fields and NN potentials, achieving substantial improvements in sampling scale and training efficiency. To facilitate reproducibility and community advancement, we open-source a high-difficulty computational chemistry benchmark.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision

Application Category

• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We introduce Adjoint Sampling, a highly scalable and efficient algorithm for learning diffusion processes that sample from unnormalized densities, or energy functions. It is the first on-policy approach that allows significantly more gradient updates than the number of energy evaluations and model samples, allowing us to scale to much larger problem settings than previously explored by similar methods. Our framework is theoretically grounded in stochastic optimal control and shares the same theoretical guarantees as Adjoint Matching, being able to train without the need for corrective measures that push samples towards the target distribution. We show how to incorporate key symmetries, as well as periodic boundary conditions, for modeling molecules in both cartesian and torsional coordinates. We demonstrate the effectiveness of our approach through extensive experiments on classical energy functions, and further scale up to neural network-based energy models where we perform amortized conformer generation across many molecular systems. To encourage further research in developing highly scalable sampling methods, we plan to open source these challenging benchmarks, where successful methods can directly impact progress in computational chemistry.

Problem

• Scalable algorithm for sampling unnormalized densities
• First on-policy approach with efficient gradient updates
• Incorporates symmetries for molecular modeling in chemistry

Innovation

• Adjoint Sampling for scalable diffusion processes
• On-policy approach with more gradient updates
• Incorporates symmetries and periodic boundary conditions