PaperlayerView exact revision on arXiv

Source-linked AI summary

Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching

Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrich, Anuroop Sriram, Brandon Wood, Daniel Levine, Bin Hu, Brandon Amos, Brian Karrer, Xiang Fu, Guan-Horng Liu, Ricky T. Q. Chen

arXiv:2504.11713v3cs.LGcs.AI

TL;DR

Sampling high-dimensional target distributions from expensive unnormalized energy functions remains difficult, especially when ground-truth samples are unavailable. Adjoint Sampling learns diffusion-based samplers through stochastic control, enabling many updates per energy evaluation and scaling to molecular conformer generation. The framework also supports molecular symmetries and periodic boundaries, with author-reported strong performance on synthetic energies and larger-scale conformer generation.

  • Problem

    Sampling from high-dimensional distributions using only unnormalized, potentially expensive energy functions is challenging, while diffusion-based generative models often require unavailable ground-truth samples.

  • Method

    Adjoint Sampling is a stochastic-control variational-inference framework using Reciprocal Adjoint Matching and a replay buffer to train diffusion samplers with many gradient updates and few energy evaluations.

  • Results

    The method achieves state-of-the-art performance on synthetic energy functions and is the first to scale to much more difficult conformer generation.

  • Takeaways & Limitations

    Adjoint Sampling supports scalable molecular sampling by incorporating symmetries and periodic boundary conditions and by enabling amortized conformer generation across molecular systems.

  • Takeaways & Limitations

    In practice, the algorithm uses a replay buffer containing samples from multiple prior steps and does not perform one optimization step until convergence, choices reported to smooth optimization and improve efficiency.

Abstract

from arXiv · show

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.

1 Introduction

Sampling from high-dimensional Boltzmann distributions using only expensive unnormalized energy functions is challenging. Adjoint Sampling addresses this with a scalable, theoretically grounded diffusion-sampling framework that supports molecular structure and amortized benchmarks.

  • Motivation: High-dimensional energy-based sampling is difficult because energy functions may be expensive and target distributions can have intricate energy landscapes.Traditional MCMC and SMC methods can also suffer from slow mixing and poor scalability.
  • Motivation: Diffusion and flow-based generative models typically require ground-truth data, which is often unavailable for molecular simulations and physics-based inference.Prior approaches therefore often add sequential Monte Carlo or importance sampling, increasing energy-function evaluations.
  • Method: Adjoint Sampling is a stochastic-control variational-inference framework for learning diffusion-based samplers from unnormalized energy functions at larger scale.It specializes and improves Adjoint Matching for efficient sampling.
  • Method: It is the first on-policy approach to allow far more gradient updates per model sample and energy evaluation, supporting difficult amortized settings.This reduces the frequency of the two costly operations required by earlier Adjoint Matching: controlled-process simulation and terminal-cost evaluation.
  • Applications and benchmarks: The framework incorporates graph and Euclidean symmetries together with periodic boundary conditions for molecular and torsional representations.The paper also introduces amortized molecule-sampling benchmarks designed to challenge methods at scale.

2 Preliminaries

The paper formulates diffusion sampling as a stochastic optimal-control problem whose controlled process should match a target distribution at a fixed terminal time. Adjoint Sampling then focuses on optimizing this control objective efficiently.

  • Controlled diffusion: The controlled SDE uses a learnable drift u and noise function σ to define time-marginal distributions over the diffusion process.The desired terminal condition is p^u_1(X_1) = µ(X_1).
  • Target path distribution: The base process sets u ≡ 0, while controlled trajectories are compared with a Schrödinger bridge connecting the initial Dirac distribution to the target distribution.Among processes matching the target at t = 1, this bridge deviates least from the base process in KL divergence.
  • Optimization objective: The sampling problem becomes a minimum-energy stochastic optimal-control problem with a terminal cost determined by the target density and base process.The normalization constant does not affect the optimal solution.
  • Optimization objective: A unique optimal control u* can achieve the target terminal distribution, reducing the task to optimizing the stochastic-control objective efficiently.Adjoint Sampling is designed to address this optimization problem.

3 Adjoint Sampling

Adjoint Sampling reformulates diffusion-based sampling from unnormalized energy functions as a stochastic-control problem and accelerates optimization by replacing repeated trajectory simulation and energy evaluation with reciprocal projections, buffered targets, and cheap updates.

  • 3 Adjoint Sampling: Adjoint Sampling solves the stochastic-control formulation of sampling by optimizing an objective whose optimum corresponds to the target distribution.The method is designed for target distributions specified only through unnormalized energy functions.
  • 3 Adjoint Sampling: The simplified Adjoint Matching formulation has a constant lean adjoint state, eliminating its additional simulation and yielding a cheaper regression loss.This follows from setting the base drift to zero, so the adjoint state remains equal to the terminal-cost gradient.
  • 3.2 Reciprocal Adjoint Matching: Reciprocal Adjoint Matching samples terminal states from the controlled process and intermediate states from closed-form base-process posteriors, producing conditionally independent time samples.The reciprocal projection replaces the unavailable controlled-process pair sampling with a Schrödinger-bridge construction sharing the terminal distribution.
  • 3.2 Reciprocal Adjoint Matching: Adjoint Sampling alternates between constructing a buffer of terminal samples and cost gradients and optimizing the control on that buffer.The regression target and terminal-state distribution are updated less frequently, enabling multiple control updates per expensive operation.
  • 3.2 Reciprocal Adjoint Matching: Closed-form Gaussian base and posterior distributions make control updates cheap and allow many gradient steps without controlled-process simulation or energy evaluation.The method therefore separates inexpensive regression updates from expensive sampling and terminal-cost computations.
  • 3.3 Adjoint Sampling Theory: The alternating scheme is theoretically an implicit projection combined with Adjoint Matching, and its fixed point is the optimal control.In practice, the authors use replay from multiple prior steps and stop short of fully converging each control update to improve smoothing and efficiency.

4 Related Work

Prior diffusion and stochastic-control samplers address unnormalized-density sampling but often depend on auxiliary sampling, expensive computation, or restrictive data assumptions. Adjoint Sampling is positioned as an on-policy alternative for computationally expensive energy functions and molecular conformer generation.

  • MCMC- and SMC-augmented learning methods improve proposals but retain scalability challenges in high-dimensional sampling.
  • Diffusion samplers for unnormalized targets often require auxiliary mechanisms to obtain training signals, including sequential Monte Carlo or importance sampling.
  • Stochastic-control samplers formulate sampling as optimization of controlled diffusion processes matching the target distribution, but remain computationally demanding.
  • Adjoint Sampling applies its diffusion-based framework to molecular conformer generation through graph-conditioned sampling.
  • Off-policy methods avoid current-model samples, whereas Adjoint Sampling is on-policy and explicitly uses energy gradients in a setting with expensive energy functions.
  • The paper identifies a gap in deep-learning methods that generate molecular conformers directly from energy without data at amortized scale.

5 Experiments

Experiments evaluate Adjoint Sampling on synthetic multi-particle energies and neural-network molecular energies, including Cartesian and torsional conformer generation. The method matches strong particle-distance performance while reducing energy evaluations and improves recall over RDKit across molecular datasets.

  • 5.1 Synthetic Energy Functions: Experiments cover DW-4, LJ-13, and LJ-55 synthetic energy functions, comparing Adjoint Sampling with iDEM, PIS, DDS, LogVariance, and an ablation without Reciprocal Projection.
  • 5.1 Synthetic Energy Functions: Evaluation reports particle W2 distance, energy-distribution W2 distance, path effective sample size, and energy evaluations per gradient update.
  • 5.1 Synthetic Energy Functions: Adjoint Sampling is on par with iDEM on particle W2 distances, while iDEM has much larger energy-distribution W2 distances on LJ-13 and LJ-55.
  • 5.1 Synthetic Energy Functions: O(100, 000) more energy evaluations per gradient update are used by iDEM than by Adjoint Sampling, while single-sample iDEM produces unstable metrics.
  • 5.2 Sampling Conformers from an Energy Function: The molecular benchmark uses eSEN energy predictions and graph-conditioned sampling in Cartesian coordinates or torsion angles for amortized conformer generation.
  • 5.2 Sampling Conformers from an Energy Function: Across SPICE and GEOM-DRUGS, Cartesian and torsional Adjoint Sampling outperform RDKit in recall, while GEOM-DRUGS precision is worse than RDKit.
  • 5.2 Sampling Conformers from an Energy Function: Without relaxation or pretraining, torsional sampling performs better than Cartesian sampling; after pretraining and relaxation, Cartesian sampling surpasses torsional sampling.
  • 5.2 Sampling Conformers from an Energy Function: The performance gap over RDKit increases as the number of rotatable bonds—and therefore conformational difficulty—increases.

6 Conclusion

Adjoint Sampling is a scalable diffusion-based sampler for unnormalized energy functions, grounded in stochastic optimal control and designed to reduce the computational burden of training.

  • Adjoint Sampling learns diffusion-based samplers from energy functions using a highly scalable approach.
  • The stochastic-control formulation connects the target sampling problem to a Schrödinger bridge with a unique optimal control under the stated conditions.
  • Adjoint Matching solves the control problem through a lean-adjoint fixed-point formulation without differentiating through the sampling procedure.
  • In its general formulation, Adjoint Matching requires separate forward trajectory and backward lean-adjoint simulations.
  • For the zero-drift, zero-terminal-cost base process, the lean-adjoint dynamics vanish, simplifying Adjoint Matching substantially.

B Base process derivations

The base process uses analytically tractable diffusion transitions, including wrapped-Gaussian constructions on periodic spaces, while reciprocal projection supplies the relevant KL-optimal path representation.

  • The base process is a controlled stochastic differential equation whose zero-control version defines the reference diffusion.
  • Forward, marginal, and backward transition distributions are derived for the base process, enabling tractable reciprocal objectives when backward kernels are available in closed form.
  • Constant and geometric noise schedules yield corresponding terminal and backward transition distributions for implementation.
  • On the torus, simulating the real-valued base process and reducing modulo 1 produces a wrapped Gaussian distribution, with dimensions sampled independently for a factorized base process.
  • Reciprocal projection maps a path measure into the reciprocal class of the reference process and minimizes KL divergence within that class.
  • The drift-level projection provides the framework for analyzing the RAM loss and its relationship to stochastic-control objectives.

C.2 Adjoint Sampling Preserves Critical Points of Adjoint Matching

Theoretical analysis shows that Adjoint Sampling preserves the critical points of Adjoint Matching, while practical training uses replay buffers and partial optimization rather than idealized exact updates.

  • Adjoint Sampling is formulated as sequential minimization of the RAM loss using samples from successive controlled processes.
  • The first-order stationarity condition for the RAM-related objective is expressed through the control and a conditional terminal-gradient expectation.
  • The critical points of Adjoint Matching coincide with the unique optimal control for the stochastic-control problem.
  • Theorem C.3 states that the Adjoint Sampling iteration is equivalent to an update preserving the relevant Adjoint Matching fixed-point structure.
  • A drift is a fixed point of the Adjoint Sampling iteration and its reciprocal projection exactly when it is a critical point of the Adjoint Matching loss.
  • In practice, finite replay buffers and partial RAM optimization differ from the idealized iterations analyzed theoretically.

E.3 Architectures and Hyper parameters

The experiments use symmetry-aware neural architectures and metrics for synthetic energies and molecular conformer generation, including Cartesian and torsional representations with periodic geometry.

  • Architectures and Hyper parameters: Synthetic-energy experiments use EGNN models with task-specific depths, hidden features, replay-buffer sizes, sample counts, and geometric noise schedules.DW-4 uses 3 layers and 128 hidden features; LJ-13 and LJ-55 use 5 layers and 128 hidden features.
  • Reported Metrics: Geometric W2 accounts for rotational and permutation symmetries when comparing generated point clouds with ground-truth samples.
  • Reported Metrics: The exact symmetry-minimized squared distance is computationally infeasible, so the experiments approximate it with a sequential search.
  • Reported Metrics: Energy W2 compares the one-dimensional energy distribution of generated samples with that from long-run MCMC simulations.
  • Reported Metrics: Path-ESS uses normalized importance weights over sampled paths and takes values in [0, 1].
  • Molecular conformers: Molecular conformers are stable local minima, and larger molecules generally have more conformers as rotatable-bond counts increase.
  • Torsion angles: Torsion angles are modeled on a flat torus with identified boundaries, reducing dimensionality while respecting their periodic geometry.
  • Molecular structure: The bond-structure regularizer penalizes violations of empirical interatomic distance bounds and is zero on the desired structure class under accurate radii.

F.4 Hyperparameters and Architecture Details for SPICE and GEOM-DRUGS

The SPICE and GEOM-DRUGS experiments use distinct equivariant neural architectures for Cartesian and torsional Adjoint Sampling, trained with replay buffers across multiple GPUs.

  • Cartesian Variant: Cartesian Adjoint Sampling uses a 12-layer EGNN with hidden feature dimension 128.Training runs for 5000 outer-loop iterations with replay-buffer sampling and distributed molecule generation across 8 GPUs.
  • Torsional Variant: Torsional Adjoint Sampling uses a 6-layer e3nn network with 32 scalar and 8 vector features, plus a pseudotorque output layer.The architecture predicts pseudoscalar torque quantities for torsional modeling.
  • Training Setup: Both variants generate new molecules and energy evaluations for distributed replay buffers, with each GPU maintaining a buffer capped at 64000 samples.The torsional setup trains for 3000 outer-loop iterations, while both setups sample batches from their replay buffers.

G Data Preparation for the Conformation Benchmark

The conformation benchmark organizes topology-only training inputs and reference conformations from SPICE and GEOM-DRUGS to evaluate amortized molecular sampling.

  • Benchmark Construction: Adjoint Sampling takes SMILES strings as inputs and outputs molecular conformations.Reference conformations were produced, relaxed, and deduplicated using RDKit, CREST, and ORCA procedures similar to GEOM-DRUGS.
  • Benchmark Construction: The released benchmark contains a training split, two molecular test splits, and reference conformations for evaluation.The benchmark is intended to foster scalable, amortized sampling algorithms.
  • SPICE Split: The SPICE training split contains 24,477 topology-only SMILES strings spanning molecules with 0 to 18 rotatable bonds.The SMILES strings define molecular topology without atomic coordinate information.
  • Test Splits: The SPICE test split contains 80 SMILES strings and 44,448 DFT-annotated, geometry-optimized conformations, while the GEOM-DRUGS test split contains 80 SMILES strings and 7,024 conformations.Both test splits are grouped by molecules with 3 to 10 rotatable bonds.
  • Benchmark Contribution: The benchmark contribution is its organization and the computational effort required to generate 44,448 SPICE conformations using a GEOM-DRUGS-like procedure.The authors describe the dataset as repackaged SPICE and GEOM-DRUGS data with additional processing for the SPICE test set.
  • Evaluation Scope: The benchmark uses recovery of released conformers as a coverage-oriented evaluation rather than a comprehensive Boltzmann-distribution sample.The released structures are local minima and are intended to represent many low-energy conformations covering configuration space.

G.1 SPICE Dataset

The SPICE preparation pipeline selects topology-defined molecules, stratifies them by rotatable bonds, generates conformers with RDKit and CREST, and refines test conformations with DFT.

  • Dataset Scope: The SPICE PubChem subset contains small drug-like molecules with 18 to 50 atoms and specified elements, while models target rotational degrees of freedom.Rotatable bonds and torsional degrees of freedom are treated as two related descriptions of molecular flexibility.
  • Flexibility Definitions: Rotatable bonds are computed with RDKit, while torsional degrees of freedom are identified using a SMARTS query and GetSubstructMatches.The torsion model receives torsional degrees of freedom, although molecules are grouped by rotatable-bond count.
  • Test Sampling: The 80-molecule SPICE test set samples 10 molecules from each rotatable-bond class from 3 through 10, while simpler 0-to-2 classes are excluded from testing.The test split is disjoint from the training split, and the procedure follows GEOM-DRUGS.
  • Conformer Generation: RDKit first generates up to 50 pruned conformers, after which the lowest-xTB-energy conformer is passed to CREST with a 6 kcal/mol final-conformer cutoff.SMILES strings encoding cis/trans stereochemistry are skipped because their strings are not preserved by the procedure.
  • Refinement and Deduplication: CREST conformers are further geometry-optimized with ORCA at a high DFT theory level before duplicate removal using a 0.1 Å RMS threshold.The stated DFT settings include wB97M-V/def2-tzvpd and tight SCF convergence controls.
  • Dataset Comparison: The SPICE test set yields an order of magnitude more conformers than GEOM-DRUGS because its molecules are more flexible and have flatter potential energy surfaces.The difference is associated with fewer rings and hydrogen bonds in SPICE molecules under the shared 6 kcal/mol CREST cutoff.
  • Flexibility Distributions: Figure 7 compares rotatable bonds with torsional degrees of freedom and uses conditional-probability colors to show their relationship.The torsional count is always greater than or equal to the rotatable-bond count.

G.2 GEOM-DRUGS Dataset

The GEOM-DRUGS evaluation uses a filtered, stratified molecular test set and its reference conformers, while energy histograms assess synthetic-benchmark sampling against MCMC.

  • GEOM-DRUGS Selection: The GEOM-DRUGS test set is filtered by element type and samples 10 molecules for each rotatable-bond count from 3 through 10.Its molecules range from 19 to 65 atoms, making the set slightly larger than SPICE by molecule size.
  • Reference Conformations: GEOM-DRUGS conformers serve as the reference for evaluation after being represented as RDKit molecules and exported as Cartesian-coordinate XYZ files.Average reference-conformer counts by rotatable-bond number are reported in Figure 8.
  • Reference Conformations: Figure 8 compares average reference-conformer counts against rotatable bonds for SPICE and GEOM-DRUGS using different y-axes.The figure separates the two datasets into left and right panels.
  • Molecular Size: Figure 9 reports SPICE atom-count frequencies for total atoms and heavy atoms, while the model predicts positions for all atoms.The two panels distinguish total atoms from hydrogen-excluding heavy atoms.
  • Synthetic Energy Evaluation: Adjoint Sampling more effectively avoids high-energy regions than ground-truth MCMC comparisons in synthetic benchmarks, especially for the higher-dimensional LJ55 surface.The comparison is shown through energy histograms for DW4, LJ13, and LJ55.

H.2 Runtime Analysis of Molecular Conformer Generation

Adjoint Sampling’s runtime analysis shows that it addresses the main computational bottlenecks in molecular conformer generation, enabling substantially more gradient updates within the same execution time.

  • Runtime Analysis: The runtime analysis specifically evaluates computational efficiency on the molecular conformer generation task.
  • Runtime Analysis: Adjoint Sampling mitigates the intensive SDE simulation and costly energy evaluation bottlenecks in prior conformer-generation approaches.These are identified as the two significant computational bottlenecks in the runtime analysis.
  • Runtime Analysis: It performs substantially more gradient updates within the same execution time, improving performance when energy-function evaluations are expensive.Conformer generation is given as an example of a task with expensive energy evaluation.
  • Experimental Setting: The eSEN energy model used in the experiments has approximately 900k trainable parameters and is configured with two layers, 64 channels, and a 4.5 Å radius cutoff.The model uses Lmax = 2 and Mmax = 0 and was reported to predict DFT energy and forces accurately on SPICE.
  • Ablation: Reciprocal Projection produces significant improvements across all reported conformer-generation metrics, especially on the unseen GEOM-DRUGS dataset.The comparison uses recall and precision metrics with coverage thresholds of 1.25Å, without pretraining or post-generation relaxation.
  • Conformer Complexity: The coverage-recall gap between Adjoint Sampling and RDKit increases for molecules with more rotatable bonds.
Loading 2504.11713v3…