Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Havens, Aaron, Miller, Benjamin Kurt, Yan, Bing, Domingo-Enrich, Carles, Sriram, Anuroop, Wood, Brandon, Levine, Daniel, Hu, Bin, Amos, Brandon, Karrer, Brian, Fu, Xiang, Liu, Guan-Horng, Chen, Ricky T. Q.
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866913863119667200
author Havens, Aaron
Miller, Benjamin Kurt
Yan, Bing
Domingo-Enrich, Carles
Sriram, Anuroop
Wood, Brandon
Levine, Daniel
Hu, Bin
Amos, Brandon
Karrer, Brian
Fu, Xiang
Liu, Guan-Horng
Chen, Ricky T. Q.
author_facet Havens, Aaron
Miller, Benjamin Kurt
Yan, Bing
Domingo-Enrich, Carles
Sriram, Anuroop
Wood, Brandon
Levine, Daniel
Hu, Bin
Amos, Brandon
Karrer, Brian
Fu, Xiang
Liu, Guan-Horng
Chen, Ricky T. Q.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2504_11713
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching
Havens, Aaron
Miller, Benjamin Kurt
Yan, Bing
Domingo-Enrich, Carles
Sriram, Anuroop
Wood, Brandon
Levine, Daniel
Hu, Bin
Amos, Brandon
Karrer, Brian
Fu, Xiang
Liu, Guan-Horng
Chen, Ricky T. Q.
Machine Learning
Artificial Intelligence
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.
title Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2504.11713