Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching
Authors: Aaron J Havens
Benjamin Kurt Miller
Bing Yan
Carles Domingo-Enrich
Anuroop Sriram
Daniel S. Levine
Brandon M Wood
Bin Hu
Brandon Amos
Brian Karrer
Xiang Fu
Guan-Horng Liu
Ricky T. Q. Chen
Venue: ICML 2025 poster (2025)
Keywords: Sampling
Stochastic Optimal Control
Primary area: probabilistic_methods->monte_carlo_and_sampling_methods
TL;DR: We introduce a highly scalable algorithm for learning to sample from only energy functions, the first of its kind in terms of efficiency, which we scale to new benchmarks on conformer generation.

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. Code \& and benchmarks provided at https://github.com/facebookresearch/adjoint_sampling.

Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching
Aaron Ha
Benjamin Kurt Miller Bing Yan", Carles Domingo-Enrich" Anuroop Sriram 2 Bin Hu Brandon Amos? Brian Karrer Xiang Fu Guan-Horng Liu"
Daniel Levine?
Brandon Wood
Ricky T. Q. Chen
Abstract
1. Introduction
Sampling from complex, eigh-dimensional distributions un- derlies many important problems in computational science, with applications spanning molecular modeling, Bayesian inference and generative modeling. In particular, we are interested in sampling from the target distribution with only access to its unnormalized energy function E, which defines the Boltzmann distribution
We introduce Adjoint Sampling, a highly scalable and efficient algorithm for learning diffusion processes that sample from unnormalized densities, 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 theoretically grounded in stochastic optimal control and shares the same theoretical guarantees 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 mod- eling molecules in both cartesian and torsional coordinates. We demonstrate the effectiveness of our approach through extensive experiments classical energy functions, and further scale 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 open source these challenging benchmarks, where successful meth- ods can directly impact progress in computational chemistry. Code & and Benchmarks provided at githubbcom/ wtthh..com/aacebooookrco facebook facebookresearchh saant_ssapping..
exp (- [(xx)
(1)
Z
where Z = Seu exp (-)) do < oo is the unknown normalization constant. The Boltzmann distribution de- scribes the equilibrium state of many physical systems, where E(x) denotes the energy of configuration æ, and 0 is a temperature parameter. Efficiently sampling from such distributions remains challenging, especially for high-dimensional systems with intricate energy landscapes. Additionally many energy functions are extremely computa- tionally expensive, e.g. requiring physics simulations.
Traditional approaches, such Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) using well- designed Markov Chains (Neal, 2001; Neal et al., 201 I; Del Moral et al., 2006), provide asymptotically unbiased samples but often suffer from slow mixing and poor seal- ability to high- nigh-dimeesional settings. This necessitates the design of better transition densities and smarter proposal dis- tributions, Recent works try to address this by augmenting sampling with learned proposal distribution (Albergo et al., 2019; Arbel et al., 2021: Gabrié et al., 2022) via normaliz- ing flows (Chen et al.. 2018; Rezende & Mohamed, 2015).
It may seem natural to look towards the recent explosion of diffusion and flow-based generative models (Song Ermon, 2019; Ho al., 2020; Lipman et al., 2023; Albergo et al., 2023). However, naïve adaptation of these data-driven generative modeling frameworks requires access ground truth data. This limitation is particularly significant in ap- plications such as molecular simulations and physics-based inference, where direct access to samples from the target dis- tribution is often unavailable. As result, prior attempts of- ten require an augmentation with sequential Monte Carlo or
'Core Contributor, Work done during internship at FAIR University of Illinois Urbana-Champaige "FAIR at Meta 'New York University 'Microsoft Research New England. Correspon- dence to: Aaron Havens <ahavens2 @ idu, Ricky T.Q. Chen <rtqichen mett.coom>>
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Adjoint Sampling: Highly-Scalabbe Diffusion Samplers via Adjoint Matching
importance-sampling (Phillips et al., 2024; De Bortoli et al., 2024; Akhound-Sadegh al., 2024), making these methods highly inefficient in terms of energy function evaluations.
2. Preliminaries
In this section, we briefly introduce an optimization problem over a class of stochastic processes that samples from p at the optimal solution the fixed time = 1. This control perspective of sample generation was observed by Tzen & Raginsky (2019b), with similar formulations found in Zhang & Chen (2022); Berner et al. (2023). A more in- depth introduction to stochastic control for sampling can be found in Appendix A.I.
The connection between sampling and diffusion processes was established by Tzen & Raginsky (2019b) through using classical results of stochastic optimal control (SOC) and Schrodinger-Briee problems (Pavon, 1989; Dai Pra. 1991: Föllmer, 2005; Chen et al., 2016). Using this formulation, Zhang& & Chen (2022) show that by directly parameterizing the drift of the controlled process, one could solve an SOC problem given unnormalized density for sampling tasks. This concept is further generalize by Berner et al. (2023) and Richter& & Berner (2024); however, all of these method require computationally expensive simulation of the diffu- sion process per gradient update. Furthermore, they require at least one sometimees many manyy evaluations per gra- dient update.
2.1. Optimizing Diffusion Processes for Sampling
Consider the stochastic differential equation (SDE)
dx, = (t) Xt, t) dt + (t)) dB, Xo= =0 (2)
where a [0. 1] + R is a scalar noise function, u : Rd x [0, 1] - R2d learnable control and (B:), is a d-dimensional Brownian motion. Under mild conditions on and u, this controlled process (2) uniquely defines time- marginals p (X,) for all E IR and we seek a drift such that pi (X,) = =(()). We denote X = (X, 0 < / < 1} a sample trajectory and p*(X) and p the distribu- tions over trajectories' generated by the controlled pro- cess (2) and base process dX,, = =(()) dB, (i.e. setting u 0 in (2)), respectively. The drift that transports the Dirac distribution the target density μ time / = is not unique: however, we can choose particular target path density that is Schrädinger bridge (Pavon, 1989; Chen et al., 2016) between the Dirac and the target distribution given by
To overcome these challenges, we introduce Adjoint Sam- pling, novel and extremely efficient variational inference framework based on stochastic control of diffusion pro- cesses, which we apply at much larger scale than previous methods. Our method is built on top of Adjoint Matching (Domingo-Enril Enrich et al., 2024), a recent method developed for solving general stochastic control problems which we specialize and improve for efficiently learning to sample.
Our contributions are as follows:
Our method is the first on-policy approach to allow far more gradient updates per model sample and energy evaluation. This is extremely important for scaling up to difficult amortized settings.
Efficiency:
p p'(x)= = huse (x (X))),(X)),
(3)
where p (X Xi) is the posterior distribution of the base process conditioned on arriving at X₁. This Schrödinger bridge is, among other processes that coincides with l at = 1, the one that deviates the least from the base process in KL divergence. Notably, the path KL to p² over trajectories with respect to the control u can be expressed analytically
Theoretically Grounded: Our method is grounded in stochastic control, following the recent work of Adjoint Matching (Domingo-Farich et al., 2024). We improve upon Adjoint Matching by proposing a new objective that implicitly projects the model onto a set of optimal controls. We do not require corrective measures such as importance sampling or sequential Monte Carlo for our theoretical guarantees.
by factorizing p" (X) p buse p p (X,) and invoking
((X))
Girsanov's Theorem (Protter & Protter, 2005).
(KL (x)) (X))
(4)
= =DxL(x(X)) II p hase (X)) + Ep log p) hase (X1)
(5)
m((x))
Structure: Graph and Fuclidean symmetries can be casily incorporated. We also adopt the method to han- dle periodic boundary conditions for modeling torsion angle representations of molecular systems.
Epp [[ (dlll(Xx,))) t) log
buse phave((X) (X1)
(6)
((X1)
To this end, we propose to minimize the objective (6), which, intuitively, minimizes the control energy required to trans- port p. In fact, this criteria corresponds exactly with a
New Benchmarks: We introduce amortized molecule sampling benchmarks that challenge new methods to be applicable at scale. Being successful at these bench- marks directly drives progress in chemistry.
Formally, path measures of stochastic processes and their KL. divergences are defined by Radon-Nikodys derivative, egg. (X) = exp(ʃ] 1 blluel] ad + Jo ue - dB..) with X ~ P..
2

---

Adjoint Sampling: Higghy--Sallbbl Diffusion Samplers via Adjoint Matching
minimum-energy stochastic optimal control (SOC) problem
Matching loss is
(x) with terminal cost function given by log [(xr This partic-
2 ///11 (x(xx,()+(()) a(t; x1| di
1
ular SOC problem is heavily studied in literature and has
LaM(u)= = E Exap'
unique optimal solution u' satisfying (Kappen, 2005; Tzen & Raginsky, 2019a) (See Appendix A. for more details)
d a(t: x) = =((; (V((x)) tb
d
s.t.
df
(((1))
a(1;X)= Vg()))
(KL("" "1p"(0 (x) (X)) = - log Hypne
O.
hase (X1)
(10)
(7)
where ā is referred to as the lean adjoint state, and stopgrad(u) denotes a stop gradient operation of u, i.e. although X ~ pd is sampled according to the controlled process.
Therefore p" can indeed be achieved by u" and our prob- lem is then reduced to optimizing the SOC objective (6) as efficiently possible. This is precisely what our proposed method Adjoint Sampling addresses.
Our first key observation is that the original Adjoint Match- ing algorithm can be significantly simplified in the case of our sampling formulation (8). In particular if we set set = 0, the lean adjoint state equation a a X) 0, X) = (X())) admits the unique analytical solu- tion a(t: X) Vx( for all t E [0, This leads to drastically simplified regression loss requiring no additional simulation of the lean adjoint state.
3. Adjoint Sampling
Now that we have equated the sampling problem as the solution of an SOC problem, we will now propose highly efficient method to solve SOC problems of the form
min UxinEx-u Ex~p U. +l
(8a)
(1 + + (XVV)(XX) dt
(uM(u)= = =Ex=pi
s.t. dX2 - (())((X,, t) di + o(1) dB,, Xo = 0 (8b)
(11)
Roughly speaking, Adjoint Matching offers a simple inter- pretation:
which we will use to learn a stochastic process that samples from a target distribution given only an unnormalized energy function (1) with
for each intermediate state X,, simply regress the control onto the negative gradient of the terminal cost
- (X,) for all possible X1 that can be arrived at from X₂. Since this is a moving target, this will, over the course of optimization, slowly shift the process towards regions with small terminal cost g-
g(x)= = log P1
buse (1) + BB()),
(9)
which is (6) without the normalization constant as it does not affect the optimal solution. Naively, one could simulate the SDE (8b), and then differentiate through the objective (8a), often referred to as the adjoint method (Bryson& & Ho, 1969), and is the method of choice for prior works (Chen et al., 2018; Zhang& & Chen, 2022). However, this is very slow as it requires two simulations per iteration, one for the state X, and another one for backpropagating through the SDE.
3.2. Reciprocal Adjoint Matching
Adjoint Matching still requires two computationally expen- sive operations at every iteration: (i) simulation of the con- trolled process and (ii) evaluation of the terminal cost. Sim- ulation of an SDE is slow iterative procedure, and the terminal cost can also be costly. For instance, in our formu- lation, the terminal cost involves the energy model F which in many real-world scenarios is extremely computationally expensive. Ideally, neither of these two operations should be done frequently. In this section, we propose modifications to the Adjoint Matching algorithm to make it highly scalable, removing both requirements.
3.1. Adjoint Matching
Adjoint Matching (Domingo-Entee Enrich et al., 2024) is an algo- rithm designed to solve a more general family of stochastic optimal control problems, being able to make use of base processes with arbitrary drift d.X, = b(X,, t) d (t) dB,, Unlike standard adjoint methods (Bryson & Ho, 1969) that use gradient- based approaches and differentiate through the objective (8a), Adjoint Matching turns the problem into a moving regression formulation. In particular, it solves for the fixed point u such that u(z, t) = (t) J x, t) where J is the value function. We discuss Adjoint Match- ing in detail in Appendix A. For our objective, the Adjoint
Firstly, notice that the training loss only depends on sam- pling (X,, ~ p", not the full trajectory. Hence (11) is equivalent to sampling pairs (i) from the joint distri- bution defined by the stochastic process p":
(AM(U)= = =/,- Ep..
[
3 [ll((x,)))) + (XX()(( di.
(12)
3

---

Adjoint Sampling: Hiighy--Sallbble Diffusion Samplers via Adjoint Matching
This does not yet yield any efficiency gains as the only way to sample X1) p" is through simulating the controlled process. Our key insight is to make use of the knowledge that at the optimal solution u" we have that pži (X,, X1) Pejl base (X₂) X1 (X) due to (3). Although the current control u does not necessarily satisfy this prop- erty, we can project the path measure- -nouurree-kwwn as a Recip-
Algorithm 1 Adjoint Sampling
I: Input: Terminal Cost: g = log pogp!ae p + / F,, base pro- cess: pi given by d(X; = ( (t))d outer-loop batch size: n, inner-loop batch size m. SDE drift network: uo. Replay buffer B.
2: B < 2
3: while Outer-Loop do
rocal projection which we make formal in Section 3.3 onto the Schrödinger bridge that generates pf given by Poli (X.]X ()) (X1). Based on this, we propose the Re- ciprocal Adjoint Matching (RAM) objective:
4: # Euler-Maruyame with no gradient
{x J==1 ~ pf. -= stopgzad(ua)
6: # gradient of energy is evaluated once per sample
7: Vg(9) ()(
B t- BU [U ()
[RAM(")= 1."
8:
9:
while Inner-Loop do
A(t)) (l + (( dt. Pappeng 2
(13)
((x+), Vg(0)) ~ ((B)
E
10:
(6)) u(0, 1)), x(5)
hose (3r) (x))
11:
That is, we sample X1 according to the controlled process, then sample X, conditioned on X1 using the posterior distri- bution defined by the base process. Since the base process is an SDE with zero drift (equivalent to (2) with = 0), con- ditionals Pij (X, I)) are known in closed form and can be easily sampled. This also has the effect of de-correlating samples across time because the X₁ samples for different values of t are now conditionally independent given X1, whereas in (12) the X, are sampled from the same trajec- tory. We additionally apply a time scaling of X(t) . which does not affect the optimal solution but improves numerical stability.
ck2 + 20 | us (X (1) .(6) + Vg)
12:
+ opoopimmzer.steppp str..ee(f.. Vo E, cl
13:
14:
end while
15: end while
16: Output: SDE sampler drift ug.
2. Obtain updated control u,+1 by optimizing (13) using samples {x1²) ( ~ B.
We theoretically justify this alternating scheme in the next section as implicitly performing a projection onto a more optimal control that aids in decreasing the SOC objective. Note that optimizing (13) in Step 2 above is extremely cheap due to the analytical form of the forward base process. Both place and the posterior distribution plyee can be designed as closed-form Gaussian distributions which we detail in Appendix B. This means that many gradient updates can be carried out without needing to simulate the controlled process or evaluation of the energy model. We provide detailed pseudo-code in Algorithm 1.
We note that (13) is related to the training objectives that appear in PDDS (Phillips et al., 2024) and TSM (De Bortoli et al., 2024), where the same formula appears inside the expectation; however, the expectations differ and makes a significant difference. In particular, PDDS and TSM take an expectation with respect to the optimal control, whereas we take an expectation with respect to the current control. That is, they require sampling from the target distribution, which is intractable. Unlike PDDS and TSM, which is simple regression formulation with a fixed regression target, our formulation has a moving regression target and is designed to solve for a fixed point (see Appendix A).
Geometric Extensions Symmetries are essential to effi- ciently sample from energy-based models defining physi- cal systems. We reduce dimensionality and improve data efficiency by enforcing symmetries in our model of the controlled process using Equivariant Graph Neural Net- works (EGNNs) (Satorras et al., 2021) or Tensor Field Net- works (Thomas et al., 2018), implemented in e e3nn (Geiger et al., 2022; Jing et al., 2022). See Appendix D for more details.
We can further increase computational efficiency by noting that we can fix the regression target and delay updating it (which we justify theoretically in Section 3.3. To design highly efficient algorithm, we further decouple p(X1) from the regression problem of learning u, delaying updates to (()) and performing multiple iterations to train u. Practically, this leads to the following alternating algorithm which we refer to as Adjoint Sampling:
3.3. Adjoint Sampling Theory
1. Using the current control uy construct buffer B (x(() ()( with samples (x std (X1) and Vg(0 ())
In this section, we provide theoretical justification to our proposed Algorithm 1. Firstly, let us define a projection P of a control onto the SOC solution where pt the target
4

---

Adjoint Sampling: Hiighy--Sallable Diffusion Samplers via Adjoint Matching
distribution:
3.4. Geometric Extensions
P(u)- = arg min (x (X) X || basc (X (X))).
Symmetries are essential to efficiently sample from energy- based models defining physical systems. Molecules are symmetric with respect to atom permutations, rotations, translations, and parity (reflection). We reduce dimension- ality and improve data efficiency by enforcing symmetries in our model of the controlled process. Our model is ei- ther parameterized by E(n) Equivariant Graph Neural Net- works (EGNNs) (Satorras et al., 2021) or Tensor Field Net- works (Thomas al., 2018), implemented in e3nn (Geiger et al., 2022; Jing et al., 2022), which incorporate graph fea- tures while respecting symmetry constraints (Kondor et al., 2018; Weiler et al., 2018: Miller et al., 2020: Geiger & Smidt, 2022). We experiment with two controlled processes: one sampling atomic positions using EGNN, and the other sampling torsion angles with e3nn. See Appendix G for details.
V
(14)
This is a projection step onto the Schrödinger bridge that shares the same terminal distribution p (X(). In connection to existing literature, this can be understood as a combina- tion of the Reciprocal and Markovian projections of Shi et al. (2023). In terms of minimizing the objective, the resulting control process consistently as effective as, if not superior to, the original control. We formalize this in the following proposition:
Proposition 3.1. After projection (14), Reciprocal Adjoins Matching is equivalent to Adjoint Matching.
(usu( LAM(((()), and furthermore, this projection improves upon on the SOC objective: J(u)= 2 (((())), where
(15)
((u):=EA [ ! 2 tt o(xi) (16)
4. Related Work
Learning Augmented MCMC Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) meth- ods have been the standard for sampling from complex distributions using a well-designed Markov-chain. Due to prohibitively long mixing times and poor scaling in high- dimensions, existing work have combined MCMC and SMC techniques with deep learning. Albergo et al. (2019), Arbel et al. (2021) and Gabrié et al. (2022) learn better MCMC proposal distributions with variational inference via normal- izing flows (Chen et al., 2018). Matthews et al. (2022) pro- posed learning proposal distributions for improving SMC using stochastic normalizing flow (Wu et al., 2020). Albergo & Vanden-Eijnder (2024) and Holderrieth et al. (2025) starts from MCMC procedures and optimize a learnable compo- nent by minimizing the Kolmogorov equations.
We defer the proof to Appendix C.I. The result of Propo- sition 3.1 hints that if we explicitly perform this projection (14) at every iteration, then we are effectively performing Adjoint Matching iterations with a more optimal control. However, performing the projection requires additional com- putation steps, e.g., using algorithm such as Bridge Match- ing (Shi et al., 2023). which we want to avoid. Fortunately, turns out that using just the RAM loss is sufficient and will implicitly include this exact projection.
Our main theoretical result is the observation that the alter- nating scheme, where the distribution of X₁ samples are fixed (step 1) and we update the control by fully converging the RAM loss (step 2), is implicitly performing the projec- tion (14) while performing Adjoint Matching on the SOC objective (6). This correspondence is stated informally in the following Theorem.
MCMC-reliant Diffusion Samplers There are many works that learn diffusion processes but rely on an aux- iliary sampling mechanism to obtain the correct signal for training. More recent advances leverage score-based dif- fusion models (Song & Ermon, 2019; Ho et al., 2020) in unnormalized sampling tasks to improve sample efficiency and scalability. Phillips et al. (2024) and De Bortoli et al. (2024) demonstrate the effectiveness of simple regression objectives, such as score matching, while Akhound-Sadee Sadegh et al. (2024) propose iterated Denoising Energy Matching (iDEM), an offline, yet biased, algorithm for learning the score from replay buffer, thereby overcoming scalability issues of simulation. Phillips et al. (2024) and De Bortoli et al. (2024) use similar loss to ours but relies on sampling from the target distribution, using sequential Monte Carlo (SMC) to obtain samples for training. In general, these ap- proaches either require access to ground truth samples via
Theorem 3.2 (Theoretical guarantees of Adjoint Sampling (informal)). Starting with any control u, performing steps I and 2 to obtain Wi+1 equivalently satisfies
SLAM (P(u:)), 8u
Ui++1 = PP(u))
(17)
8LAu denotes the functional derivative with respect where bu to the control u. Moreover, the fixed point where PP())
and u = -u-u-
5Cau (u) is the optimal control u² to (8). 0u
A more precise statement and proof can be found in Ap- pendix C. This result provides the theoretical justification for our proposed algorithm (Algorithm 1). In practice, we differ slightly and use a replay buffer that contains samples from multiple prior steps, and we do not perform step 2 until convergence. We find this helps smoothen the optimization and improve computational efficiency.
5

---

Adjoint Sampling: Heighly-Sallable Diffusion Samplers via Adjoint Matching
Table 1. Results for the synthetic energy function experiments. We report geometric W2 metric hased on Klein al. (2024h) and ID energy histogram E() W2 metric (visualized in Appendix 1.1) with respect to ground truth MCMC samples, We also report a path- measure ESS when applicable. See Appendix F.4 for more details. 'The values reported are per sample (i.e., divided by the batch size) and according to the LJ-55 experiment hyperparameters.
DW-4 (d 8)
LJ-13 (d m 39)
LJ-55 (d 165)
evals evals if evals evals
per gradious per gradient updlate"
Method
path ESS 1
W/ 4
E(-) W/ 4
path ESS| 1
W/ 4
E(-) W/ 4
path-ESS+ ESS 1
W/3 4
E(-) Wg 4
updlate'
PIS kowg
0,462
0.68
0.65
0.012
1,93
18.02
0.001
4.79
22
1000
DDS Vergo aut,
0,461 s3e
0.92
0.90
0.010-am
1.99
24.6
0.001
4.60
173.09
1000
Log Varianco (ldtarn & bemer 2015
0.025e:sc
1.04
188
1000
iDEM Salghht 2004
a.70
0.55
1.61 1o01
30.78_2400
4.69
512
3
iDEM w/ MC sample
1.21 Loor
2.70
2o03
22.41
5.79
1e32
Adjoint Sampling w/o RP (Ablation)
0.448 g4448
0.63g.n
1.03
0.159
1.68 1001
0.094
94.48
0.002
Adjoint Sampling (Ours)
0.627 m0.0
0.62 m0r
0.55
0.220 m354
1.67
2,40
0.066
4.50
58.04-28.5
0.002
an auxiliary sampling algorithm such as MCMC, or rely importance- weiightee estimation or resampling, all of which require extensive evaluation of the energy function and it is not clear whether these improve upon simple data-driven learning algorithms (Appendix E).
et al., 2024) or prescribed (Albergo & Vanden-Eijnden, 2024; Chen et al., 2024a; Holderrieth et al., 2025). In con- trast, Adjoint Sampling is an on-policy training method, explicitly uses the gradient of the energy function, but bene- fits significantly from requiring only X1) pairs, being able to use a replay buffer and Reciprocal projections to train at significantly reduced computational cost.
SOC-based Diffusion Samplers In contrast, stochastic optimal control (SOC) based samplers reframe sampling tasks as an optimization problem. Zhang & Chen (2022) and Vargas et al. (2023) showcase directly optimizing for controlled processes that match the desired target distribu- tion. This framework has been generalized in many ways Berner et al. (2023); Richter& & Berner (2024); Vargas al. (2024); Chen et al. (2024a). However, all of these methods are hindered by their computational requirements, including computationally expensive differentiation through the sam- pling procedure, computation of higher-order derivatives in constructing the training objectives, or the need for impor- tance sampling (i.e. multiple energy evaluations). Adjoint Sampling overcomes these challenges by being able make significantly more gradient updates per generated sample and energy evaluation.
Molecule Conformer Generative Models To our knowl- edge, deep learning- -ased sampling method has been applied at scale to generate molecular conformers directly from energy without data. Boltzmann generators use im- portance sampling to improve performance using energy information, but require offline molecular dynamics data to learn the initial models (Noé et al., 2019). Using an an- nealed importance sampling method, Boltzmann generators have been able to to sample small, single particle systems without data, but their results are not amortized and require many energy evaluations (Midgley et al., 2024). Several flow- generative modeling works have learned to gen- crate conformers and Boltzmann distributions by training on ground truth data (Köhler et al., 2020; Jing et al., 2022; Xu et al., 2022; Hassan al., 2024; Köhler et al., 2023; Klein et al., 2024a; Diez al., 2024). These works did not learn solely from existing energy models.
Off-policy Methods Off-policy methods (Malkin et al., 2023; Richter& & Berner, 2023; Achound-Sadege et al., 2024; Hua et al., 2024) are those which do not need to use samples from the current model. As such, off-policy methods do not inherently make use of the gradient of the energy function, as they do not differentiate through the model sample. To alleviate this, it is typical to parameterize the model using the gradient of the energy function, and it has been found that performance strongly relies on this trick (He et al., 2025). As our setting concerns computationally expensive energy functions, this parameterization is no longer feasi- ble. Furthermore, being off-policy does not imply being more efficient. Many off- methods still require the full trajectory to evaluate their loss, only being able to use sub- cub-trajectories if the time amme-marginals are either additionally learned (Bengio et al., 2023; Lahlou et al., 2023; Zhang
5. Experiments
We now evaluate Adjoint Sampling on several energy func- tions over multi-particle systems. We compare to previous works (Zhang & Chen, 2022; Akhound-Sadegh et al., 2024) synthetic energy benchmarks, whose energy is known analytically and are very cheap to evaluate. We then propose challenging molecular conformer generation benchmark, where one trains amortized models to sample conformers for a large dataset of organic molecules. For this task, we used the eSEN energy model (Fu al., 2025).
6

---

Adjoint Sampling:
Hiigly--Sallbbl Diffussin Samplers via Adjoint Matching
Reciprocal Projection ablation We include an ablation of Adjoint Sampling, where instead of using the Reciprocal projection (RP) to sample Xt given X1, we simply store sample pairs (X1,X) in the buffer and train on the Adjoint Matching objective (12). We call this Adjoint Sampling w/o RP, which helps demonstrate the effectiveness of the Recip- rocal projection. Ablation results for conformer generation are differed Appendix 1.4 Table 5.
in Appendix 1.1). Moreover, iDEM uses (100, 000) more energy evaluations than Adjoint Sampling per gradient up- date since it relies on Monte Carlo (MC) estimation of the ground truth score function. This is what makes scaling iDEM to more complex energy functions intractable. To test a more efficient version of iDEM, we also run iDEM us- ing only a single MC sample for estimating their objective, however it proved to yield overall unstable metrics, most likely due to its biased MC score target estimator. In regards to SOC-based methods, PIS and DDS struggle to perform beyond IJ13, while the offline method Log Variance can not scale beyond DW-4. These results indicate that Adjoint Sampling is the preferred method for large and complex energy models.
5.1. Synthetic Energy Functions
Several synthetic energy functions of n-particle bodies pro- posed and benchmarked by previous works (Köhler al., 2020; Midgley et al., 2023; Klein al., 2024b; Akhound- Sadegh et al., 2024). We consider three energy functions: A 2D 4-particle Double- Well Potential (DW-4), 3D 13- particle Lennard- potential (LJ-13) and finally a 55- particle Lennard-Jone energy (LJ-55). See Appendix F for details.
5.2. Sampling Conformers from an Energy Function
We provide a new benchmark to encourage research into highly scalable sampling methods capable of finding con- formers from a molecular energy model. These are local minima the molecule's potential energy surface, i.e. fixed molecular topology has set of conformers determined by locally stable configurations of rotations around its bonds. Depending on the flexibility, a molecule can have one or many different conformers.
Baselines Along with the Adjoint Sampling ablation, we compare against the most recent ofte-of-the-ar sampler iDEM (Akhound-Sadeg) et al., 2024) who claims to be first learned sampler to scale successfully to LJ-55. We also compare against PIS (Zhang & Chen, 2022) and DDS (Var- gas et al., 2023), being the most comparable method to Adjoint Sampling in its SOC / reverse KL formulation. Lastly, for completeness we present an offline variant of the PIS loss referred to as the LogVariance (Richter & Berner, 2023), which is actually gradient-free free. All methods use the EGNN (Satorras et al., 2021) architecture for all experiment settings, where PIS and DDS use fewer number of layers making baoo-propo through the SDE tractable.
Sampling conformations based on quantum mechanical en- ergy functions such as density functional theory (DFT) is important but challenging task. Due to the high cost of DFT calculations, we use energy model that was pre- trained on the SPICE-MACE-OFE dataset (Eastman et al., 2023; Kovács et al., 2023) called eSEN (Fu et al., 2025). It has shown very high accuracy in predicting the DFT en- ergy. eSEN is a transfcrable message passing neural network model that can predict the energy of given molecular con- formation for molecules containing (H. C, N. O, F, P. S, CI, Br, 1). It featurizes molecule with sparse graph represen- tation, encoding the chemical element in the node features and the Cartesian displacements in the edge features. The SPICE-MACF-ON dataset comprehensively samples the energy surface in high and low energy regions, enabling eSEN to be generally applicable to many conformations of small molecules.
Evaluation In Table 1 we report our metrics including Wy: the geometric 2- Wasserstien distance (taking into ac- count symmetries by first optimizing a rigid transformation (Klein et al., 2024b) between model generated samples and the ground truth MCMC data and E(:)Wy The ID W/2 distance between the distribution of energy values produced by generated samples and ground-truth and path effective sample size (ESS): Roughly measuring the proportion of informative sample trajectories with respect to each algo- rithms target SDE process. Since iDEM is not built on a reverse KL objective, the path-ESS metric is not meaningful and is hence omitted. Please see Appendix F.4 for more details. To demonstrate the energy sample efficiency of Adjoint Sampling we report the number of energy function evaluations per gradient update.
Using the Adjoint Sampling framework, we can sample atomic coordinates of a specific molecule by conditioning on its connectivity graph. These atomic coordinates are proposed conformations. The graph conditional method is appealing as enables amortization, i.e. learning a sam- pler for molecules consisting of arbitrary combinations of organic elements and bonds. We design two approaches:
Results It can be seen clearly in Table 1 that Adjoint Sampling is on par with iDEM in terms of particle W-2 distances, however iDEM struggle to remain in low energy regions as indicated by the much larger E()) on the LJ- 13 and LJ-55 experiments (visualized as energy histograms
1. Cartesian Adjoint Sampling (Appendix G.I): We sample all 3D Cartesian coordinates for each atom, con- ditioned on elements (nodes) and bond type (edges).
7

---

Adjoint Sampling: Higghy--Sallbbl Diffusion Samplers via Adjoint Matching
Table 2. Recall and precision metrics for large scale amortized conformer generation. Coverage values are for thresholds of 1 1.25A.
Standard deviations are computed across molecules in the test set. Results split by number of rotatable bonds can found in Figure 10.
SPICE
GEOM-DRUGS
Recall
Precision
Recall
Precision
Method
Cov.
AMR 4
Cov. f
AMR
Cov. f
AMR 4
Cov.
AMR 4
RDKi ETKDG
72.74
1.04.052
69.68+37.11
1.14
63.51
1.15 #0.i
69.77=
1.09
Torsional AdjSampling
85.06
0.86
70.42
1.06
0.98
67.85 17.00
1.09
Cartesian AdjSampling
82.22
0.96=022
49.13+3300
1.26
60,93
1.20
28.44
1.86
Cartesian AdjSampling (+pretrain)
8944
0.84
65.93
1.12
72.98 #:30:2
1.02
45.14
1.47
relaxation
RDKi ETKDG
81.61 4:2158
0.79
74.51 4:3507
0.97
71.72+22.73
0.93
75.37
0.89
Torsional AdjSampling
88.25 12
0.72
74.66
0.94
76.62
0.87 10:00
71.85 123.54
0.97
Cartesian AdjSampling
94.10
0.68 0022
62.80
1.02
79.08 4:9.44
0.89 20:45
47.88
30.92
1.00nn
Cartesian AdjSumpling (+pretrain)
96.65 7.51
0.60
77.20
0.92 0.92+
87.01 2:22.9
0.76 =..34
61.96-34
1.10
Cartesian
Torsional
Xa
X1
Generation process
Figure 1. The figure depicts two sampled trajectories from trained Adjoint Sampling models that use either cartesian or the torsional representations. They larget conformations of the held-out SMILES string COCSc1sc2cccce2 [n+]] [0-] The left frame Xo comes from the initial Dirac distribution and right frame X, is sampled conformer.
2. Torsional Adjoint Sampling (Appendix G.2): We sample only the torsion angles in the molecule, con- ditioned on elements, bond lengths, and bond angles (3-muless).
learning. All our models are trained on SPICE; however, we utilize merely the SMILES strings to extract the molec- ular topology. This is possible because Adjoint Sampling does not require any data during training, i.e., we do not use the atomic configurations! We evaluate generalization capabilities by computing precision and recall on subset of GEOM-DRUGS, see Appendix H.2. This is one the largest datasets and most relevant to pharmaceutical drug design.
In addition, we evaluate a pretraining step on RDKit sam- ples using Bridge Matching (Shi et al., 2023) which aids in initial exploration of low energy regions (see Appendix E scale the eSEN
explor
for more details). In experiments, energy by the Boltzmann constant and a low temperature. In Cartesian Adjoint Sampling, we introduce additional energy terms to regularize bond lengths, ensuring we are target- ing the correct molecule according to its connectivity. The terms do not appear in Torsional Adjoint Sampling; they are unnecessary since the torsional SDE cannot break bonds by design.
We split the SPICE molecules into a train set and a 80 molecule test set, allowing us to validate our sampler's ability to extrapolate to unseen molecules. We produce reference conformers for evaluation purposes using RDKit (Landrum, 2013); CREST (Pracht et al., 2020), which is considered the gold standard in molecular chemistry; and ORCA (Neese, 2012). Because this procedure is computational expensive, we do this only for the SPICE test set. GEOM-DRUGS provides conformers for us. See Appendix for details about the datasets.
Datasets We consider two datasets of molecular struc- tures: SPICE (Eastman et al., 2023) and GEOM- DRUGS (Axelrod & GomezzBombarelll, 2022). SPICE enumerates over 23,000 diverse drug-like molecules and their properties; it is designed to be utilized for machine
Baselines We compare our approach to RDKit ETKDG (Riniker & Landrum, 2015), a chemistry- conformer
8

---

Adjoint Sampling: Higgly-Scallbbl Diffusion Samplers via Adjoint Matching
SPICE
HEOM-DRUGS
relaxation n1 relaxation
relaxation
Coverage Recall (%)
Coverage Recall (%)
R1 reloxation
80
60
40
40
20
Torsional AdjSampling
20
Torsional AdjSampling
Camesian ad[Sampling
Camessin Ad|Samplin
Camesian adjSamping, pretrain)
Cartestan Ad[Sampling pretrain)
Threshold (Angström)
Threshold (Angström)
Figure 2. Recall coverage versus RMSD threshold for Adjoint Sampling variants and RDKit. show performance both with and without relaxation.
generation method. RDKit combines rules-based distance matrix method, iterative refinement algorithm, 3D em- bedding, and experimental torsional preferences to generate conformers, making it suitable domain-specific baseline.
by being initialized as noisy stochastic process.
Without relaxation or pretraining, the Torsional AdjSam- pling variant performs significantly better than Cartesian Variant in precision and recall on both datasets. However, after pretraining Cartesian AdjSampling surpasses the Tor- sional variant and even further exceeds in performance after relaxation, (being pretrained or not). This may be due the Torsional representation evolving on a constrained space as seen in Figure 1, where the Cartesian representation evolves with the energy regularizer (see Appendix G.1) but is un- constrained. This state representation will affect how the controlled process explores the configuration space, suggest- ing that the unconstrained Cartesian representation samples are more diverse.
Evaluation Our assessment focuses on reference con- former coverage recall, which measures the percentage of reference CREST conformers recovered by the gener- ated samples, and coverage precision, which evaluates how often generated conformers closely match low-energy ref- erence structure. In particular, recall is very important metric for molecular design because exploration is typically challenging and good mechanisms exist to filter candidates. The number of reference conformers varies significantly for molecules under consideration, ranging from tens to thou- sands. We also report Absolute Mean RMSD (AMR) as an error metrie, with lower values indicating better struc- tural fidelity. See Appendix G.3 for more details. Results are presented both with and without relaxation, where post- generation optimization helps refine molecular conformers in Table 2. While samples without relaxation are a test of the sampling algorithm's 's performance, testing with re- laxation is ultimately what brings the most value com- putational chemistry being standard refinement process (see Appendix G).
Finally, in Figure 10 also compare across different num- ber of rotatable tonds-the number of conformers increases drastically with the number of rotatable bonds (see Fig- ure where we see that the gap between our model and RDKit increases as the difficulty increases.
6. Conclusion
We introduce Adjoint Sampling, highly scalable approach for learning diffusion-based samplers from energy functions. Using the Reciprocal Adjoint Matching objective and a replay buffer, our method enables efficient training by allowing many gradient updates with few energy evaluations and model samples. Based on continuous-time diffusion processes, our framework naturally integrates symmetries and periodic boundary conditions, making it effective for molecular modeling and conformer generation. We achieve state-of- sitee-of-th-aare art performance on synthetic energy functions and are the first to scale up to much more difficult conformer generation. We open-source our benchmarks to encourage further advancements highly scalable sampling methods.
Results As seen in Table 2, across both SPICE and GEOM-DRUGS, Adjoint Sampling (both Cartesian and Torsional) outperform the RDKit baseline in Recall. On generalizing to the GEOM-DRUGS dataset, we show worse precision than RDKit but significantly improved recall, As depicted in Figure 2, Cartesian Adjoint Sampling outper- forms RDKit almost all threshold values, with and with- out relaxation. While relaxation improves coverage across all methods, it particularly benefits the performance of Ad- joint Sampling, making the gap between RDKi--ETDGG even larger. We believe this comes from the fact that Adjoint Sampling naturally explores more of the configuration space

---

Adjoint Sampling: Hiighy--callbbl Diffusion Samplers via Adjoint Matching
Impact Statement
Berner, J., Richter, L.,, and Ullrich, K. An optimal con- trol perspective diffusion-basee generative modeling. arXiv preprint ar.iv:2211.0135555 2023.
Our method primarily serves scientific and industrial re- search and does not pose immediate risks for misuse. How- ever, with any advancement in generative modeling, there is broader responsibility ensure its use aligns with ethi- cal guidelines, particularly in fields like drug discovery and materials engineering, where biases in models or training data could impact real-world outcomes. Our work has di- rect upplications in computational chemistry, particularly in molecular modeling and conformer generation, which could accelerate scientific discovery in drug design, material sci- ence, and biophysics.
Bryson, A. E. and Ho, C. Applied Optimal Control. Blaisdell, 1969.
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled langevin diffusions. arXiv preprint arXiv:24/2.0708/ 2024a.
Chen, R. T. Q.. Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
References
Chen, Y., Georgiou, T., and Pavon, M. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:671-691, 2016.
Akhound-Sadegh, T., Rector-Brooks, J., Bose, J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., et al. Iterated denoising en- ergy matching for sampling from boltzmann densities. In Forty-first International Conference on Machine Learn- ing, 2024.
Chen, Y., Goldstein, M., Hua, M., Albergo, M. S., Boffi, M., and Canden-Fijnded with stochastic interpolants and follmer processes. arXiv preprint arXiv:2403.13724. 2024b.
Probabilistic forecasting
Albergo, M. S. and Vanden-Fijnden, E. Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711, 2024.
Dai Pra, P. A stochastic control approach to reciprocal dif- fusion processes. Applied mathematics and Optimization, 23(1):313-22) 1991.
Albergo, M. S., Kanwar, G., and Shanahan, P. E. Flow- based generative models for markov chain monte carlo in lattice field theory. Physical Review D. 103(3)003555, 2019.
De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., and Doucet, A. Riemannian score- bosee-baee generative modelling. Advances in Neural Information Processing Systems, 35:2406-2440 2422, 2022.
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E.
Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint 27X1v:2303.08997, 2023.
De Bortoli, V., Hutchinson, M., Wirnsbegee, P., and Doucet, A. Target score matching. arXiv preprint arXiv:2402.088 086668 2024.
Arbel, M., Matthews, A., and Doucet, Annealed flow transport monte carlo. In International Conference on Machine Learning, pp. 318-338830333 -30... PMLR, 2021.
Del Moral, P., Doucet, A., and Jasra, A. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411 (133:411-436 1-436, 2006.
Axelrod, S. and Bomez--oobarelli, R. Geom, energy- annotated molecular conformations for property predic- tion and molecular generation. Scientific Data, 9(1):185, 2022.
Diez, J. V., Atance, S. R., Engkvist, O., and Olsson, S. Gen- eration of conformational ensembles small molecules via surrogate model-assisted molecular dynamics. Ma- chine Learning: Science and Technology, 5(2):025010. 2024.
Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., and Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Ad- vances in neural information processing systems, 35: 11423-11436, 2022.
Domingo-Enrich, C., Drozdzal, M., Karrer, B., and Chen, R. T. Q. Adjoint matching: Fine-tuning flow and diffu- sion generative models with memoryless stochastic opti- mal control, 2024. URL https: http:::s///xxvvoor//// abs/
Bellman, R. Dynamic programming. Princeton Landmarks in Mathematics. Princeton University Press, Princeton, NJ, 2010., 1957.
Bengio, Y., Lahlou, S., Deleu, T., Hu, J., Tiwari, M., and Bengio, E. Gflownet foundations. Journal of Machine Learning Research, 24(210):1-55, 2023.
Eastman, P., Behara, P. K., Dotson, L.,, Galvelis, R., Herr, J. E., Horton, J. T., Mao, Y., Chodera, J. D., Pritchard, P., Wang, Y., et al. Spice, a dataset of drug-like
10

---

Adjoint Sampling: Highly-Scalable Diffusion Samplers via Adjoint Matching
molecules and peptides for training machine learning potentials. Scientific Data, 10(1):11, 2023.
He, J., Du, Y., Vargas, F., Zhang, D., Padhy, S., Yunang, R., Gomes, C., and Hernández-Looato, J. M. No trick, no treat: Pursuits and challenges towards simulation- free training of neural samplers. arXiv preprint arXiv:2502.0668. 2025.
Fleming, W. and Rishel, R. Deterministic and Stochas- tic Optimal Control. Stochastic Modelling and Applied Probability. Springer New York, 2012.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, volume 33. Curran Associates, Inc., 2020.
Föllmer, H. An entropy approach to the time reversal of diffusion processes. In Stochastie Differential Systems Filtering and Control: Proceedings of the IFIP-WG 7/1 Working Conference Marseille-Luminn France, March 12-17, 1984, pp. 156-163. Springer, 2005.
Holderrieth, P., Albergo, M. S., and Jaakkola, T. LEAPS: A discrete neural sampler via locally equivariant networks. arXiv preprint arXiv:2502.10844 2025.
Friede, M., Hölzer, C., Ehlert, S., and Grimme, S. dxtb-an efficient and fully differentiable framework for extended tight-binding. The Journal of Chemical Physics, (6)(6), 2024.
Hua, M., Laurière, M., and Vanden-Eijnden, simulation-free deep learning approach to stochastic opti- mal control. arXiv preprint arXiv:2410.05100 2024.
Fu, X., Wood, B. M., Barroso-Luque, I.., Levine, D. S., Gao, M., Dzamba, M., and Zitnick, C. L. Learning smooth and expressive interatomic potentials for physical property prediction. arXiv preprint arXiv:2502.12147 2025.
Huang, C.-W., Aghajohari, M., Bose, J., Panangaden, P., and Courville, A. C. Riemannian diffusion models. Ad- vances in Neural Information Processing Systems, 35: 2750-2761, 2022.
Gabrié, M., Rotskoff, G. M., and Vanden- Eijnnenn, E. Adap- tive monte carlo augmented with normalizing flows. Pro- ceedings of the National Academy of Sciences, 119(10), mar 2022.
Jing. B., Corso, G., Chang. J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 24240 22253, 2022.
Ganea, O., Pattanaik, L., Coley, C., Barzilay, R., Jensen, K., Green, W., and Jaakkola, T. Geomol: Torsional geo- metric generation molecular 3d conformer ensembles. Advances in Neural Information Processing Systems, 34: 13757-13769, 2021.
Kappen, H. J. Path integrals and symmetry breaking for optimal control theory. Journal of Statistical Mechanics: Theory and Experiment, 2005(11) 1), nov 2005.
Geiger, M. and Smidt, T. e3nn: Euclidean neural networks, 2022. URL httpa::/arxxv..og///aa 2207.09453.
Karras, T. Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion brafusioo-bbsed generative models. Advances in neural information processing systems, 35: 26565-26577, 2022.
Geiger, M., Smidt, T., M., A.. Miller, B. K., Boomsma, W., Dice, B., Lapchevskyi, K., Weiler, M., Tyszkiewicz, M., Batzner, S., Madisetti, D., Uhrin, M., Frellsen, J., Jung, N., Sanborn, S., Wen, M., Rackers, J., Rød, M., and Bailey, M. Euclidean neural networks: e3nn, April 2022. URL https://dol.rrg/....5. denodo.645933
Klein, L., Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., Noé, F,, and Tomioka, Timewarp: Transferable acceleration of molecular dy- namics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024a.
Klein, L.,, Krämer, A., and Noé, F. Equivariant flow match- ing. Advances in Neural Information Processing Systems, 36, 2024b.
Grimme, S., Bannwarth, C., and Shushkov, P. A robust and accurate tight-binding quantum chemical method for structures, vibrational frequencies, and noncovalent interactions of large molecular systems parametrized for all spd-block elements (z= 1-86). Journal of chemical theory and computation, 3(5)::989-2009) 2017.
Köhler, J., Klein, L.., and Noé, Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference machine learning, pp. 5361- 5370. PMLR, 2020.
Hassan, M., Shenoy, N., Lee, J., Stark, H., Thaler, S., and Beaini, D. Et-flow: Equivariant flow-matching for molec- ular conformer generation. In The eiihth- Annual Conference on Neural Information Processing Systems, 2024.
Köhler, J., Invernizzi. M., De Haan, P., and Noé, F. Rigid body flows for sampling molecular crystal structures. International Conference on Machine Learning, pp. 17301-17326. PMLR, 2023.

---

Adjoint Sampling: Higgly--ealable Diffusion Samplers via Adjoint Matching
Kondor, R., Lin, Z., and Trivedi, S. Clebsch-gordan nets: a fully fourier space spherical convolutional neu- ral network, 2018. URL https: /ttp///aaxxx/////////// abs/ 1806.09231.
Miller, B. K., Geiger, M., Smidt, T. E., and Noé, F. Relevance of rotationally equivariant convolutions for predicting molecular properties. arXiv preprint arXiv:2008.n 000:2000..0461,, 2020.
Kovács, D. P., Moore, J. H., Browning, N. J., Batatia, I., Horton, J. T., Kapil, V., Witt, W. C., Magdäu, I.-B., Cole, D. J., and Csányi, G. Mace-off23: Transferable machine learning force fields for organic molecules, arXiv preprint arXiv:2312. 2023.
Neal, R. M. Annealed importance sampling. Statistics and computing, 1 11: 11:125-139, 2001.
Neal, R. M. et al. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2(11):2, 2011.
Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández-Garcet A., Ezzine, L. N., Bengio, Y., and Malkin, N. A theory of continuous generative flow net- works. In International Conference on Machine Learning. pp. 18269-18300, PMLR, 2023.
Neese, F. The orca program system. WIRES Comput, Molec. Sci., 2(1):73-78, 2012. doi: 10. 10022/cmss...
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(647)) 147, 2019.
Landrum, G. Rdkit documentation. Release, 1(1-79):4, 2013.
Pavon, M. Stochastic control and nonequilibrium thermody- namical systems. Applied Mathematics and Optimization, 1921,7-202, 1989.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Peluchetti, S. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023a.
Liu, G.-H., Vahdat, A., Huang, D.-A., Theodorou, E. A., Nie, W., and Anandkumar, A. Pssb: tonge-to-image schr\" odinger bridge. arXiv preprint arXiv:2332.05582, 2023.
Peluchetti, S. Non-denoising forward time diffusions. arXiv preprint arXiv:2312. 14589, 2023b.
Liu, G.-H., Lipman, Y., Nickel, M., Karrer, B., Theodorou, E., and Chen, R. T. Q. Generalized schrödinger bridge matching. In The Twelfth International Conference Learning Representations, 2024.
Phillips, A., Dau, H.-D., Hutchinson, M. J., De Bortoli, V., Deligiannidis, G., and Doucet, A. Particle denois- ing diffusion sampler. arXiv preprint arXiv:2402.06320, 2024.
Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers. In European Conference on Computer Vision,
Pracht, P., Bohle, F., and Grimme, S. Automated explo- ration of the low-energy chemical space with fast quan- tum chemical methods. Physical Chemistry Chemical Physics, 22(14):7169-7199 2020.
pp. 23-40. Springer, 2024.
Pracht, P., Grimme, S., Bannwarth, C., Bohle, F., Ehlert, S.. Feldmann, G., Gorges, J., Müller, M., Neudecker, T., Plett, C., et al. Crest-a program for the exploration of low-energy molecular chemical space. The Journal of Chemical Physics, 160(11), 2024.
Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory balance: Improved credit assignment in gflownets. arXiv preprint arXiv:220J. 13259, 2023.
Matthews, A., Arbel, M., Rezende, D. J., and Doucet, A. Continual repeated annealed flow transport monte carlo. In International Conference Machine Learning, pp- 15196-15219, PML... 2022.
Protter, P. E. and Protter, P. E. Stochastic differential equa- tions. Springer, 2005.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In Proceedings of the 32nd Interna- tional Conference on Machine Learning, 2015.
Midgley, L., Stimper, V., Antorán, J., Mathieu, E., Schölkopf, B., and Hernández-Loobato, J. M. SE(3) equiv- ariant augmented coupling flows. Advances in Neural Information Processing Systems, 36, 2024.
Richter, and Berner, J. Improved sampling via learned diffusions. arXiv preprint arXiv:2307.01 2023.
Midgley, L. I., Stimper, V., Simm, G. N., Schölkopf. B., and Hernández-Lobatee J. M. Flow annealed importance sam- pling bootstrap. In The Twelfth International Conference on Learning Representations: ICLR 2024, 2023.
Richter, L. and Berner, J. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations, 2024.
12

---

Adjoint Sampling: Highly-Scalable Diffusion Samplers via Adjoint Matching
Riniker, S. and Landrum, G. A. Better informed distance geometry: using what know improve conforma- tion generation. Journal of chemical information and modeling, 55(12):2562-25 25774, 2015.
Veber, D. E., Johnson, S. R., Cheng, H.-Y., Smith, B. R., Ward, K. W., and Kopple, K. D. Molecular properties that influence the oral bioavailability of drug candidates. Journal of medicinal chemistry. 45(12):2615 2623,, 2002.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference on machine learning, pp. 9323-9332. PMLR, 2021.
Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. 3d steerable cnns: Learning rotationally equivariant features in volumetric data, 2018. URL attps://arxivg
Sethi, S. Optimal Control Theory: Applications to Man- agement Science and Economics. Springer International Publishing. 2018.
Wu, H., Köhler, J., and Noé, F. Stochastic normalizing flows. Advances in Neural Information Processing Systems, 33: 5933-5944, 2020.
Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffu- sion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2023.
Xu, M., Yu, L.,, Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular con- formation generation. arXiv preprint arXiv:2203.0292 2022.
Somnath, V. R., Pariset, M., Hsieh, Y.-P., Martinez, M. R., Krause, A., and Bunne, C. Aligned diffusion schrödinger bridges. In Uncertainty in Artificial Intelligence, pp. 1985-1995. PMLR, 2023.
Zhang, D., Chen, R. T. Q., Liu, C.-H., Courville, A., and Bengio, Y. Diffusion generative flow samplers: Im- proving learning signals through partial trajectory op- timization. In The Twelfth International Conference Learning Representutions, 2024. URL https:// openreview.net conttt/oouuungn wow.mmid-oOsseeaee
Song. Y. and Ermon, S. Generative modeling by estimat- ing gradients of data distribution. arXiv preprint urXiv:/907.080 2019.
Zhang, Q. and Chen, Path integral sampler: A stochastic control approach for sampling. In International Confer- ence Learning Representations, 2022.
Song, Y., Sohk-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR 2021), 2021.
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation- and translation-qquivarann neural networks for 3d point clouds, 2018. URL https: ptttpp://arxiv.oot /axiv.org/////// abs/1 802.08219.
Thornton, J., Hutchinson, M., Mathieu, E., De Bortoli, V., Teh, Y. W., and Doucet, A. Riemannian diffusion schr)." odinger bridge. arXiv preprint wwXiv:2207.03024,, 2022.
Tzen, B. and Raginsky, M. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. 19X5v:9905,099000 2019a.
Tzen, B. and Raginsky, M. Theoretical guarantees for sam- pling and inference in generative models with latent dif- fusions. arXiv:
2019b.
Vargas, F., Grathwohl, W. S., and Doucet, A. Denoising diffusion samplers. The Eleventh International Con- ference on Learning Representations, 2023.
Vargas, F., Padhy, S., Blessing, D., and Nusken, N. Trans- port meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learning Representations: ICLR 2024, 2024.
13

---

Adjoint Sampling: Heghly-Scalable Diffusion Samplers via Adjoint Matching
A. Additional preliminaries
A.1. Stochastic optimal control
Stochastic optimal control (SOC; (Bellman, 1957; Fleming & Rishel, 2012: Sethi, 2018)) considers general optimization problems over stochastic differential equations. Specifically, a class of SOC problems can be expressed as the following optimization problem:
min / (2 011 - dt + e(H)]]
(18)
wEt/ dxp' (b(X), 1) (t))(X)) ()) dt (t)dBr
X Po
(19)
where in (19), X E I& the state of the stochastic process, Rd x [0, 1] Rd is commonly referred as the control,
b: Bd x [0, - Rd is a base drift, and g j0, 1] - add is the diffusion coefficient. These jointly define the controlled process X ~ p" that we are interested in optimizing; often both b and g are fixed and we only optimize over the control u.
In the following proposition show that for the case of b = 0, f = 0 and Xo 0 that the Schrodinge--ridge p (X)
p base
X Xx, (X) can actually be achieved by p'X)).
Proposition A.1. For the stochastic process (2), there exists a unique minimizer 4 to the following optimization problem
min DR.. (p" (p"(())))))) (x)(((())) 0,
(20)
and the optimal controlled distribution yd satisfiex:
Pt. (X,, (X,, pluse (X )x(X),
(21)
where pose (1) is the posterior distribution of the base process conditioned on X1.
Proof. By using the definition of the KL divergence of path measures, (xX =p buse (X (X(x))u(Xi) can be factorized as
((x1)
ppanee (x1)
DKL (x)lx) DKL (x"())llp hase
(X)
DKL (p"(X) hase (X)) + E Ex log
p base (X1)
((X1)
(22)
By Girsanov's Theorem (Protter & Protter, 2005) to further evaluate the right-hand side as
pX1) (x1) = Ex-p" [.' 1 alll((x,, t)] dt + log
base (X1)
p hose (X₂) p³ (X1)
DKL (p" (X) hase (X)) + Ex-p" log
This criteria corresponds exactly with minimum-energy stochastic optimal control (SOC) problem with stage cost [[((X,,))) ? and terminal cost function g(X,) = logg
( (x.))
min Exapr up 1o' 2 2ll t)l dt 9(X1))
(23)
s.t. dX, = a())a(((,))) d) o()) dB., Xo 0
(24)
We can define the cost-to-go of our problem at a particular state and time t by
Jw, x, Ex-p" [/' Alll((Xe t)1 +( X, x
(25)
then we know that (under mild conditions on g) the value function or cost-to-go under the optimal control takes rather surprising form and that indeed the optimal control u" is unique and takes the form of a time-varying function of the current state (Kappen, 2005):
V(a, t) min J[u 2,,) t) logg E,o exp( (ap(-9(Xx))) X, = =]], u'(a,t) ((t))
(26)
14

---

Adjoint Sampling: Higghy--Sallbbl Diffusion Samplers via Adjoint Matching
With this fact about the optimal solution of the SOC problem, we can upper bound the aforementioned path KL divergence between the optimally controlled process and our target Schrodinger-bridge as
DKL (p" (x)lle V(Xo: 0) log Rx (0x_ppoooxppp((
(27)
log / / (xx)dx==) -x-dd-0 =0,
(28)
By the data processing inequality we have
DKL (P xi)llet, (Xe, Xi), & DEL (x)li = 0, E [0
(29)
which implies that (X,,,))) = (X, for all E [0.1]. This is exactly our claim and completes the proof.
A.2. Adjoint Matching
Adjoint Matching (Domingo-Enrice et al., 2024) is an algorithm designed to solve stochastic optimal control problems of the form (18). Unlike standard adjoint methods (Bryson & Ho, 1969) that differentiate through the objective and perform gradient descent, Adjoint Matching directly tries to find the fixed point
u(I t) = (tt
(30)
where J is the expected future cost according to the control u,
Ep- L. (llu(X.
J(u; x, =
+) f(X., s)ds + g(X+) X1 x
(31)
Adjoint Matching then solves this fixed point (30) by replacing 2 with a stochastic estimator. This stochastic estimator is the lean adjoint (Domingo-Eniee et al., 2024) denoted by a. Domingo-Enrich et al. (2024) then showed that the unique solution to the fixed point
u(x,t) -o (t) (ā(z, () x]
(32)
the optimal control u This then motivated the following objective as a means to solve this fixed point problem.
2 .//% [lu(Xe, ot) T(t x)|| dt
1
(uM(U) == Eps
X~ pp, ā stsooggrad(w),
(33)
where
ā(t: X)= -āā(t; x) T vzb(X,, (x,(x,,)),,
(34)
dt
a(1: = (X1g(X)))
(35)
where à is referred to the "lean adjoint" state, and stopgrad(-) denotes stop gradient operation, i.c., although X pa is sampled according to the controlled process, Adjoint Matching does not differentiate through the sampling procedure.
Note that for the general problem formulation (18), Adjoint Matching requires two simulations, one to sample a trajectory from the stochastic process X ~ p" and one to solve the lean adjoint state backwards in time (34) from a terminal condition at (35).
Adjoint Matching is greatly simplified for simple base processes. One of our key observations is that the Adjoint Matching algorithm greatly simplifies when the 6 = 0 and f = 0, since then (34) is zero. This means that the lean adjoint state ä(t, X) (g(X,) for all t. The Adjoint Matching algorithm hence reduces to
11
1
/./ [lu(Xe,)) + Vxxx))
(aM(w)= = Epa
X X~p²
ū = stopgrad(u).
(36)
15

---

Adjoint Sampling: Heghly-Scallbbl Diffusion Samplers via Adjoint Matching
B. Base process derivations
B.1. Euclidean space
Our base process is modeled by a stochastic differential equation (X, E R" and a + R+);
dX, o(t)dBi,
Xu = 0
(37)
corresponding to (2) with the control zero, Le., u4
With this, the forward transition distributions
where V//= = =f. a(s)
Pils X,) = [N(x[Xs,
(38)
With 8 = 0 and Xo O, we obtain the time marginals of (37):
where = / (ss)
base (x) = -x(x00, Pt.
(39)
Furthermore, the backward transition distributions are (t > s):
plase (x) m (Xela)
hase (x))x.)
(N Dsjc()s
(40)
plo00 (X,)
Us
1
where Oult
(41)
Vs + Vels
1/8
Ueje
As such, the Reciprocal Adjoint Matching objective (13) becomes tractable when Vr|s can be derived in closed form.
Constant noise schedule. Setting the diffusion coefficient to be constant (() = a results in
M/js o² (t 8),
(42)
and hence the following terminal and backward transition distributions:
base X1) = NXX
basc Piji
(x X1)
(( t)tf).
(43)
Geometric noise schedule. in reverse-time)
Song et al. (2021); Karras et al. (2022) propose the following geometric schedule (stated here
miin
Jaax Juin
log
Baa Guis
(44)
results in
2a
Oania Triex
2+
2t
Pt|s
Crain
Ocale
21
(45)
Orui
and hence the following terminal and backward transition distributions:
base (X1) = NXx;; o
(1
Ciiin
:) I),
Paji
buse
(X,)X)) = -x oup
(46)
B.2. Flat tori
Simulating the base process (37) in = IR/Z produces a wrapped Gaussian distribution p
buse
(x: =
Pr
buse
(a
k),
(47)
kEZ
16

---

Adjoint Sampling: Highly-Scallbbl Diffusion Samplers via Adjoint Matching
where phae is the time marginal of the base process over I defined in (39) and & is the set of integers. We can equivalently can interpret (47) as marginalizing random variable k. That is, we view plase (x k) as a joint distribution (z, over E T and & E 2,
(r k) = Pt
base (a k) JET hove (x' + k))
Pt
base (a' + k))
(48)
z'eT
PP((
Intuitively, the variable k represents which sub-interval on R was collapsed to obtain æ. Furthermore, we can infer l given æ
at time d:
(x(
Pr (k|xz) = Ex Pe(ar, k')
(49)
hence arriving at the following backward transition distribution (s < t) in which we marginalize k:
pl (Xe + k)
base Palt
(x) =
Eaii(elk X.)p,(k)
Σ KEZ
Σ
Pale
hase (x (x+ (x+
EMee Pt huse (X, k²)
(50)
KEZ
K'EZ
where buse Xx + k)) is the backward transition probability in defined in (40). Thus sampling from pal (x|X1) given in (50) can done by the following sampling process
1. Sample l p(k) where p(k) x plose (X1 + k).
2. Sample Y, Pii1
pa (11x k))
3. Project X, Yg mod 1.0.
For the n-dimensional T", since the base process is factorized, we simply sample independently for each dimension.
C. Theory of Reciprocal Adjoint Matching
As a preliminary, let us first define precisely what meant be Reciprocal projection. We take some terminology and results from Shi al. (2023), denoting M the subset of Markov path measures associated with an SDE of the form dX, = v(X,, + (t) B,, with a, v locally Lipschitz.
Definition C.1 (Reciprocal class, Def. 3 of Shi et al. (2023)). A Borel probability measure IT E P(C)) is in the Reciprocal class R(Q) of Q E . if I
The Reciprocal projection of P E P(C) is defined as II" = Q (P) = Po.r (o By construction, II" belongs to Reciprocal class R(Q). The following characterization shows that the Reciprocal projection is indeed a projection, i.e.
(R( PR(O)-
Lemma C.2 (Reciprocal projection characterization, Prop. of Shi et al. (2023)). Let E P(C), II = -PRQQ) Then, II arg m (P)(QQ) E R (Q)).
Throughout the paper, we set Q = By, i.e. Q is the probability measure of the process (10 (8) dB.) dB.) which we denote as PR(Q)-
Although this projection is exactly what used by the proposed reciprocal adjoint matching objective Equation (13), we will now proceed to define similar projection in terms of drifts which will allow to analyze an idealized version of the adjoint sampling algorithm.
C.1. Reciprocal projections decrease the SOC objective
In the case that nd IP is generated by some drift function u, can view the reciprocal projection as a projection of drifts P proposed in Theorem 3.2, noting that the unique drift generating PR (p²) is given the solution of the following Schrodinger bridge problem
P(u) := arg min Ex x [//" 1 [luo(Xe; t)11 dt + log
(bbsc(())
)
(51)
(X1)
17

---

Adjoint Sampling: Hiigly-Scalabbl Diffusion Samplers via Adjoint Matching
This projection of drifts allows to reason about the RAM loss with respect both the AM and SOC functionals which described in Proposition 3.1.
Proof of Proposition 3.1. Let P(u). Because (u) P(u), the reciprocal projection "approximation" introduced in RAM is exactly the distribution induced by and so (u) = LRAM(U).
To show that ((u 2 J(P(u)), we simply apply the definition of P(u) (51) which yields
Lo 1 t)]
pluse
[(u) Ex-
log
(X11)
plase
D"
Ex-pa
1 allu(Xe, t)1? dt +
log
+ log
(X1)
Du
1
poase
log D" (x1)
2 Ex-p"
(x,, 1)11
log
(X1)
+ Exxrp"
by (51)
pd
[.¹0 211o(X,, ()1 de +
base
p"" ) (xx))
=Ex-p=
log
+ log
ppa
C.2. Adjoint Sampling Preserves Critical Points of Adjoint Matching
Consider slightly more general version of the Reciprocal Adjoint Matching (RAM) objective with arbitrary fixed reference drift u,
L(u;e) Ex~r [/- 11
( (]l((x,,)) (
2 ll((x,,t)
(52)
s.l. dX, - (()u((X,,
Xo 0
(53)
Then the RAM and AM loss can be both be equivalently restated in terms
CRAM(U)+ = C(u; ü), [aM(u) = C(u, ū) s.t. == P(ū), ū = st stopgrad(u),
where RAM is shown by first projecting via P, and then setting this the reference drift with stop gradient applied.
If in our main algorithm we consider sequentially minimizing the RAM loss where starting from some drift v; we only populate the buffer B with samples from p": we can model adjoint sampling via the following update procedure:
0;41 arg min Llu; Üs), s.t. 0 P(())
(54)
with this iterative formulation, can actually re-frame Adjoint Sampling in terms of the stationary condition of the Adjoint matching loss derived in (Domingo-Enrice Enrich et al., 2024). To show this, we first consider the first variation of L with respect to the first argument which is defined by
d
((())))))))))))))))))) (X,,)) di [f,
8L (l(wv))())
L(u emp v) Exapr de
the first-order necessary condition an extremum of this cost functional is that SL(u; w) = 0 which say
6C(u, = xx U. (u(Xe,))) (t) (V((X.) X.] mxX,, t)) dt 0, Vn E C(R x 10, 1R.
(55)
We know that if p" has support over Rd and Ep 2 (u; t) = 0 almost everywhere. This is the functional derivative with respect to which satisfies
[ [6 (X,, t). m t)) dt = 0 for all y,, then it must be that (J
8C Tu (u; = u(),)(()))))))))))))))))<>)))= + (1) (Vg(X1) Xx = 2| = 0, a.e.
(56)
18

---

Adjoint Sampling: Higgly--callbbl Diffusion Samplers via Adjoint Matching
In view of this general necessary condition, we can identify the corresponding necessary condition for (AM(").
ÜLAM (z) du
8L (u; (z, u(2, t) + (t) (Vg(Xi) X x] O, a.e.
(57)
du
This coincides with the functional derivative derived in (Domingo-Enrice et al., 2024), where it was observed that the critical points of Adjoint Matching are exactly when (57) is zero. Furthermore, it was determined that the only critical point the optimal control u² the SOC problem (8). We can now state Theorem 3.2 more precisely and provide a proof in the following Theorem.
Theorem C.3 (Adjoint Sampling). Suppose we have a feedback u, Then performing the iteration (54) to obtain u,41 is equivalent to the iteration
5LAM (P(u))),
U1+1 =P(us)
(58)
du
Moreover, u is fixed point satisfying P(u)) and u u- ( the optimal control u² to (8)).
(u) if and only if is critical point of Lam (i.e. it is
du
Proof. By minimizing L(u;ū)) point-wise for fixed i; Plu:). we know the optimal iterates take the form
12+1 = (t() p.. (V X; = z].
(59)
Recalling the Adjoint Matching loss LAM(U)= = L(u; ū) we can identify its functional derivative at P(u.) that stems from first variation to be
8LAM (P(u:)) = P((s)) +((EE (V)(Xi) X. =x] xr].
(60)
8u
Through a simple algebraic manipulation, we can arrive a new expression for the iterate given by
Ue+1 P(us) 8LAM
(P(u:)) (x, l)-
8u
We will now show that v is a fixed point of the Adjoint Sampling iteration and P(v) = e and only if o is a critical point of CaM (and hence v = u²). First of all, if o is a critical point of LAM, then it must be v = u" (Domingo-Enre) al., 2024) and 6LAM (v) O. Because v = u² the optimal solution of the SOC loss, it is its own projection and so 8u
P(u") - ACAM (P(x)))(u") = u". Therefore u is a fixed point of the Adjoint Sampling iteration.
On the other hand, if e is a fixed-point of the Adjoint Sampling iteration and P(c) = v, our iteration formula tell us that éCam (v) and so is critical point of LAM and that u² necessarily. 5u
There are some important practical difference between this statement and what we do in practice. In practice, we optimize over finite size buffer B, where the iterations discussed here in Theorem C.3 assumes we sample from p Furthermore, we do not need to exactly solve the RAM step convergence in practice and only partially optimize for fixed number of iterations before collecting new samples for B. This may be desirable as it will have a smoothing affect on our functional gradient descent of Theorem C.3 which tracks large unit step-size gradient update.
D. Geometric Extensions
Symmetries are essential efficiently sample from energy- models defining physical systems. Molecules are symmetric with respect to atom permutations, rotations, translations, and parity (reflection). We reduce dimensionality and improve data efficiency by enforcing symmetries in model of the controlled process. Our model is either parameterized by E(() Equivariant Graph Neural Networks (EGNNs) (Satorras et al., 2021) or Tensor Field Networks (Thomas et al., 201 8), implemented in eän (Geiger et al., 2022; Jing et al., 2022), which incorporate graph features while respecting symmetry constraints (Kondor et al., 2018; Weiler et al., 2018; Miller et al., 2020; Geiger & Smidt, 2022). We experiment with two controlled processes: one sampling atomic positions using EGNN, and the other sampling torsion angles with e enn. See Appendix G for details.
19

---

Adjoint Sampling: Highly-Scalable Diffusion Samplers via Adjoint Matching
D.I. Graph-Conditioned SE(3)- IIvarraat Sampling
Consider the graph ((V, &). Each node 0j V and edge eij E & has feature attributes relating the topological structure of the data (e.g., atom types and bond orderings), including spatial coordinate 2; E R.. We can condition an equivariant drift ue on the graph features such that it is equivariant to the spatial symmetries given by the group g E G == Aut(G) x SO(d) acting via = (P & Riz. (Le. uo(g X2, t: 9) g uo uo(X, (X,, G)). Here P Aut(g) are graph automorphisms represented as permutation matrix that reorders the graph nodes while preserving its structure - and It E SO(d) is an d-dimensional rotation matrix.
Translation Invariance To enforce translation invariance, the system is restricted to the zero center-of-mas. subspace X'CoM {ir - o² = 0} Rddd This is achieved by projecting the particle positions y R&dd to the zero CoM subspace using the projection operator:
1 k: 1.11
x Ay,
A = IN
2 la
(61)
where J is the identity matrix, and 1g is a vector of ones. With this we can define a zero CoM process.
dx, - o() (anAuu(Xxr:1 1; (; (B)AAd,, Xo =0
Now phase is singular Gaussian plaeee (ar) (x; mAA ), which can be sampled by first sampling the isotropic Gaussian hase also projects
Na; 0, pt) and then projecting via A. We have pr (x) (x; 0, v4) for all x E 2.CM and the posterior Pip onto CoM
Geometric Adjoint Sampling By ensuring the drift is Gequivariant and zero CoM, we ensure the model distributions are G-invariant for all E [0, 1] (Köhler et al., 2020; Xu et al., 2022). Putting this all together, our RAM loss (13) can be modified to support the proposed symmetries.
LGCORAM(0) = L'' X(1)
T (X ((Vo)))))((V(())))) dt
(62)
Assuming that energy E and slase G-invariant, the target o (X)) for learned drift will be G-equivariant and zero CoM.
D.2. Periodic boundary conditions
In many cases we may want to model a state space with periodic boundary conditions. Concretely, we consider state space that is the flat tori in n dimensions, denoted T = R"' /2". It is quotient space resulting from identifying any point - (x'....,> x², ..... xx") with (x',...,.' + +1...,, xx) for all i [m.]. The derivations for the SOC objective in R" (6) can directly be extended to this quotient space and also general Riemannian manifolds (De Bortoli et al., 2022: Thornton et al., 2022; Huang et al., 2022). Let us denote by prsse the distribution of the SDE modeling the base process, iee., (2) with n, = 0, that lives on T The RAM objective (13) requires computing V log p. (-) as part of the terminal cost, and sampling X; from phise (-IX2). Simulating the base process in T" produces factorized wrapped Gaussian distribution.
T1
Pt
(x)
Pt
(xf + k),
(63)
where plse the distribution of the base process in R. In practice, we can easily compute this up to high numerical precision by truncating the summation in (63). Furthermore, we can sample from the backwards transition distribution independently for each dimension E [n]]
h ~ p(k)),
p(k+) x p (Xi k+),
x = Yi mod 1.0,
Y Ph Paj1
buse (Y: + R²))
(64)
In practice, we consider only values of l within some truncated set [-keune--. kaune) with kannne sufficiently large to cover the region where base (Xi + k) is practically non-zero. Details for the derivation can be found in Appendix B.2.
20

---

Adjoint Sampling: Highly-SSalabbl Diffusion Samplers via Adjoint Matching
E. Bridge Matching and Pretraining
There may be situations where we have some initial data distribution plaia but want to improve our generated samples even further using an energy function. We can in fact pretrain our control drift ug on this data using Bridge Matching (Peluchetti, 2023a;b; Shi et al., 2023; Liu al., 2024; 2023; Somnath al., 2023; Ma al., 2024; Chen al., 2024b), which can be similarly formulated as controlled SDE diffusing from Xu- 0 and switch to Adjoint Sampling from the energy without changing the optimal distribution.
Specifically, given our controlled stochastic process Equation (2), we can write down the following Bridge Matching loss to regress onto
L Mt)B (x, (X (x(( (ooi))))o((t(a (o)(t) ott)² x. log Pije
buse (X1 dl,,
(65)
where for general base process dX,, = o(t) dB,, we have
Vile : (( a(s) ds.
Pile (a Xx) (x
base
Therefore our score target is
1(1 [_1(x: - + const =o(1)2 211)2x111 Xx
V log logpl Ceapphapxxx (X1]X))= =
hose
(
DIje
Vije
Additionally, apply scaling X(1) == (+)) stability, resulting in the following Bridge Matching objective.
the loss that the target is (X1 - Xt) which improves numerical
(BM(O)= = L." E8 (X. (x)xx)))
Pije
2
up (X,,t) - (X1 -x)
Xx)
dt
(66)
o(1)
F. Synthetic Energy Experiment Details
F.1. Double Well Potential (DW-4)
We the same double-well potential in iDEM (Akhound-Sadegh al., 2024), which was originally proposed in Köhler al. (2020). DW-4 describes pair-wise distance potential energy for a system of particles { x2, T3, I1), where each particle has 2 spatial dimensions E
I (d 8). The potentials analytical form is given by:
1 Σ a(de - do) + b(de) - da)2 c(di; - da) di ||x, - 22y122 = +
E(x)=
(67)
ij
where we set a = 0, 6 = -4, == 0.9 and temperature 1.
F.2 Lennard-Jones Potential (I.J-13, LJ-55)
Similarly to DW-4, the Lennard-Jones potential is also based pair-wise distances of system with n particles, but each each particle has 3 spatial dimensions. Its analytical form is given by
12
Ed (x) = E <>
I' dig
I'
d;j = ||= - =2ylla
(68)
dis
where Tm, 6 and are physical constants. As in Köhler et al. (2020) and Akhound-Sadegh et al. (2024), we add the additional a harmonic potential:
Eusc (x) -
1 Σ |1, Σ
2c
(69)
where Iсом refers to the center of mass of the system. Therefore, the final energy is then T = Ex + (x)'(()). for c the oscillator scale. As in previous work, we use I'm = 1,7 1,6 e and c 1.0.
21

---

Adjoint Sampling: Highly-Scalable Diffusion Samplers via Adjoint Matching
F33 Architectures and Hyper parameters
In this section we give brief overview of the parameters used train Adjoint Sampling on the synthetic energy functions. For all experiments, we an EGNN architecture similar to Schound-Sadepee et al. (2024). Results for iDEM on all experiments reproduced using the same configurations found in Akhound-Sadeen et al. (2024).
DW-4: For the DW-4 energy we use an EGNN with 3 layers and 128 hidden features. We train Adjoint Sampling for 1000 outer iterations, generating 512 new samples and energy evaluations per iteration into buffer of max size 10000. In each iteration we optimize on 500 batches of batch size 512 from the replay buffer, using a learning rate of x 10-4 We use a geometric noise schedule with min 10-4 and mar 3.0.
For the LJ-13 energy we an EGNN with 5 layers and 128 hidden features. We train Adjoint Sampling for 1000 outer iterations, generating 1024 new samples and energy evaluations per iteration into buffer of max size 10000. In each iteration we optimize on 500 batches from the replay buffer of batch size 512, using a learning rate of 3 x 10 We use geometric noise schedule with min = 10 and maax = 3.0.
IJ-13:
LJ-55: For the LJ-55 energy we use an EGNN with 5 layers and 128 hidden features. We train Adjoint Sampling for 1000 outer iterations, generating 128 new samples and energy evaluations per iteration into buffer of max size 10000. each iteration we optimize on 500 batches from the replay buffer of batch size 512, using a learning rate of 10 We use geometric noise schedule with miin 0.5 and 0 maz = 3.0.
F.4. Reported Metries
Geometric W2 Because our energy and diffusion process invariant rotational and permutations symmetries, we also take these symmetries into account when measuring distance between generated and ground truth point clouds (e-g., from long run MCMC). Note this W2 metric is different from what is reported in Akhound-Sadegh al. (2024), which uses the euclidean metric.
We(u, (inf / ((r,)))((, de dy
(70)
where is the transport plan with marginals constrained and respectively. Here the distance metric d takes into account all point-cloud symmetries, and is obtained by minimizing the squared Euclidean distance over all possible combinations of rotations, reflections (Od)), and permutations (S(k)) for particles of spatial dimension d.
d(zu, x1) = min
||=0- (R P)
(71)
RES())
However, computing the exact minimal squared distance is computationally infeasible in practice. Therefore, we adopt the approach of Köhler al. (2020) and approximate the minimizer by performing a sequential search
x1) R min 1R00 (R & P) p= arg min I|:O Px1113 RESO(d)
(72)
PES(K)
Energy W2 (F(-) W2) An informative way to assess the quality of the generated samples is to look at their energy distribution with respect to the ground truth distribution obtained from energies of long-run MCMC simulations. This shows how well the generated samples are avoiding high-energy regions. Specifically, we use the Wasserstein-2 (W2) distance on the I-dimensional energy distribution on R. This results in the standard euclidean W2 metric
W2(Fp, Fr)) = inf 7
g)(x (xx
T T(x
v12 dr dy
(73)
Path effective sample size (path-ESS). We use importance weights over paths X pp (Zhang & Chen, 2022),
w(X) = dp"(x)
((x))
dp dp (X)
exp (-1 (-2/// [lu(X,,)) 2t 1 t)(X),t) s(xi)).
2
(74)
(X)
1 (X1)
(X)
22

---

Adjoint Sampling: Heghly-SSallbbl Diffusion Samplers via Adjoint Matching
As these are unnormalized weights, we use the normalized effective sample size:
w!)
ESS
(75)
nE ~ (
t i=1 ii
estimated over n samples X; ~ pd with weights wi = w(X)). Additionally, this ESS is normalized by the number of samples n and takes a value in [0, 1].
sid
G. Sampling Molecular Conformers
Conformer generation is a fundamental task in computational chemistry with applications from drug discovery to catalyzing particular reactions of interest. In the context of this paper, conformers represent different stable states, i.c. local minima on the potential energy surface, of the same molecule. The number of conformers a molecule has is generally proportional to the number of rotatable bonds present in the molecule. For example, large molecules with lots of rotatable bonds have many conformers. bond is considered rotatable if there is sufficiently low energy barrier to rotation and the bond is attached to non-terminal heavy atoms (Veber al., 2002). We define a specific substructure match this definition in Appendix G.2 and use it determine the torsion angles sampled by torsion model. However, apply slightly different, more restrictive, definition to classify molecules into strata based on the standard definition from RDKit in Appendix H.
In many molecules and polymers torsion angles are particularly important degrees of freedom for sampling conformations. A molecule has 3k - 6 degrees of freedom corresponding to relevant atomic motion. This quantity subtracts global translation (3 degrees) and orientation in space (3 degrees), which are both irrelevant to the intrinsic molecule. The remaining degrees of freedom can be broken down into bond lengths (pairs of atoms), bond angles (3-tuples containing two atoms bonded to one common neighbor), and torsion angles (4-tuples containing a central bond and two atoms connected at either end). The energy required to stretch a bond or distort bond angle is typically much higher than rotating torsion angle. As result, torsion angles are the most important degrees of freedom for determining the distribution of possible molecular states, the Boltzmann distribution, and only allowing torsion angles to rotate is a common approximation used in many algorithms.
Relaxation Conformers represent local minima in the energy landscape; therefore, we are interested in descent algorithms to refine structures generated by Adjoint Sampling or RDKit to reach those minima. We apply the same energy model, eSEN, that we utilized to train our drift estimator on energy to define an optimization target. Doing gradient descent on energy using forces called relaxation. In our relaxations, we use the LBFGS algorithm and bound the maximum force norm Fmax == max; V x; E(XI)l on any of the l atoms stopping criteria, where Xi is the latest structure in the descent sequence. We stop when Fmax < 0.0154 eV/À or 200 descent iterations, whatever comes first.
G.I. Sampling the target molecule in Cartesian coordinates
Bond Structure Energy Regularization Although we condition on the molecular graph, our drift control estimator aims to approximate gradient of the energy function. During the evolution of our stochastic trajectories energy troughs that have alternative connectivity than that of our target molecule. Since we aim to generate conformers of a specific target molecule, we need to avoid this. We do so by adding a regularizer to the energy. This regularizer assigns atoms particular bond / no bond lengths based on empirically known interatomic radii and repulsive van der Waals radii.
may wander through
B(x) - - + (x)
(76)
where a > 0 is regularizer scaling constant. "flat-bottom" potential having the property that
The regularizer energy function E Rdd + [0, oo) is designed to be a
ET (z) 0 =
(77)
where S is the desired molecular structure class. In practice, we slightly relax this to be E S (xc = 0, for training stability and due to our inability to exactly characterize S.
In practice, we determine an pouer-bound distance using well-established empirical covalent radii between atoms with
bonds, and lower-bour bound distances using van der Waal radii between atoms without bonds. We then penalize any pairwise distances which violate these bounds taking into account a relaxation factor 7 2 1. With the adjacency matrix A (aij)
23

---

Adjoint Sampling: Highly-Scallbbl Diffusion Samplers via Adjoint Matching
and atom type attribute vector H (h.), determinine the molecule structure class S(A, H). With accompanying interatomic bond limit functions drona (ha ht) and ona bond ht), compute the regularizer iterating over all edges as follows:
ET (2, A,H17)) =
Σ i.jEhoo
max {1lx - 2z112 7. chond (hi, (1),,,,)))00 11
(78)
max
dlo-bonu (hy, hi) - 1| - 21112.03
(79)
A.tETwoin bood
where /omd and bomd are the index pairs of A corresponding to bond or no- between atoms respectively. Assuming accurate interatomic radiis, our regularizer satisfies x E E(AA, H) = Ereg A, H:7) = for all 2 1.
G.2. Torsion Angles and Torques
Figure 3. Multiple representations of the molecule diiodoethane with SMILES string ICCL (Left) Molecular graph representation. (Right) 3D coordinates representation. We chose the torsion angle as I-C-C-I, where indicates purple iodine atom and indicates a gray carbon atom. We identify the torsions with 4-tuple of indices, selecting the heaviest atoms on either side of the central bond as the first and last members. Ties are broken by arbitrary atomic index. We take the zero dihedral angle to be when first and last members of the 4-tuple are close, the so-called cis-isomer. Following Jing et al. (2022), we adjust the dihedral angle by rotating all the atoms at one end of the torsion about the bond axis, leaving the remaining atoms in place. This can be described as a torque pointing along the bond axis. When the torque would asymmetrically act the molecule, we define the positive direction towards the side of the molecule with more atoms.
Representing Torsion Angles and Torques The orientation around a given bond can be quantified with a torsion angle. Specifically, a torsion angle is defined between four atoms (a, b, c, d), where the torsion angle the angle between the planes formed by a, b, e and b, c, d, which rotated around the central be bond. The choice of atoms a and d are typically unique, thus different choice of d can lead a different torsion angle to describe the same molecular conformation. Our convention is that and d correspond the heaviest atoms either side of the central bond bc. We break any ties by arbitrary existing atomic index. Next, we represent torsion angles symbolically.. .nticalll...Asuming that atoms a, b, d E R are positions in Cartesian space, we define a few vectors that lie along undirected bond axes:
dr = ab b-
dx3 = ad = d
drg = bc = c b.
(80)
Those vectors allow us to define torsion angle follows:
p(dz1, dr2, dz3) = atan2 (dx2 ((drı x dzz) x (dx, x dæs)). |dxzl (dæ) x dxz) (dxz x drs)).
(81)
We define the function T: R4x3 - I that takes in four column vectors of atomic positions (a, b, c, d) and outputs a torsion angle by (81), iee., T: (a, b, c, d) + (ddx (a, b), drz(b, c), drg(e, d)).
Following Jing et al. (2022), bond axis bc. This can described as a torque in the standard angle-axis representation where the axis points along be.
adjust a torsion angle by rotating all the atoms on one side of the be bond about the
24

---

Adjoint Sampling: Heighl--Sallable Diffusion Samplers via Adjoint Matching
positive sign indicates positive rotation according the right-hand-rule and the magnitude indicates the rotation's oss effect on the existing torsion angle. The action of the rotation R(t) on our system X, with constituent atoms Iy E TR3, can be written as follows:
R(t),; ['()(,,-)) 6) + if C(c) 21
(82)
otherwise
where (t) is the rotation matrix defined by the axis-angle representation of t, and C(c) denotes the connected component of è in an identical graph to the molecule excepr with bond be removed! E(c) therefore denotes one "half." or "side," of the molecule when hypothetically divided at bond bc.
Since every torque points along be by design, we need a convention for ordering torsion angle 4-tuples to determine the direction of rotation. We choose atom e to sit on the side of the be bond with more atoms (Jing et al., 2022). To be more specific, if one broke bond bc, there would be two connected components c(b) and E(c). We choose e such that |e(c)] > |e(b)] (This the opposite convention to the definition Jing et al. (2022) write in their paper, but the same as in their code.) Given ties, we the current arbitrary index.
Finding torsional degrees of freedom with RDKit Sampling conformations by sampling torsion angles requires identifying the torsional degrees of freedom in our target molecule. We use the package RDKit (Landrum, 2013) to find and represent the 4-tuples for those degrees of freedom. We look for matches to the following substructure: (1) two atoms must be connected by bond, (2) neither end of that bond can be aromatic, (3) neither end of that bond is merely connected only to a single hydrogen, and (3) only one end of the bond can be part of ring. Such a substructure can be encoded in a SMARTS string, specifically [ [!$ (*#*) [D]] & &! (*#*) [D]]. This matching can be done on directly on the SMILES string of a molecule.
Substructure matching identifies the torsional degrees of freedom with a list of 4-tuples containing atoms in every position. Recall, this representation is not unique, as described above! Furthermore, this search technique typically finds more torsional degrees of freedom than RDKit returns when querying the number of rotatable bonds. RDKit has a specific definition of rotatable bond as defined implicitly in Ca cNumRotatak tattlaaleeent
Considerations for torsion angles geometry We model a stochastic differential equation in the torsion angle space where can reason with the base distribution more practically and the dimensionality is reduced. However, this comes with few additional considerations. For this section, consider molecule with dis torsional degrees of freedom. Due to the periodic boundary conditions on torsion angles, the space is modeled as flat torus D == [- π) ) where the extrema are identified with zero displacement due to the torus distance function, or so-called logmap. Let Dice (l,...,dre, ) denote the index set for each torsion angle.
We use a Dirac source distribution centered at y = 0 with E D at ( = 0. This choice of origin in DD must have a corresponding unique position in cartesian space. Given atoms (a,, by, c, d,) that are constituents of torsion angle ip; E [-w, т with i E Drot, we choose 9o: = 0 to be the state in Cartesian space where a, d, are closest in terms of Euclidean distance, i.e.,
T(ai, bj, Cr d,) =
d; arg min ||a; dg 2|2
(83)
The definition is the same for all i E Drot- We rely intuition and ignore the details to prove that this choice is unique, and that the distance between a, and d, is the only necessary component to specify the torsion angle yii.
We now construct the molecule Xo in cartesian space that has po; = 0 for all i E Dros- We first embed the molecular graph in cartesian space to an arbitrary, approximate conformer using RDKit's EmbedMo 1 lecule e (or EmbedMule t tiileeonfff with numConfs=1). Once embedded, we can identify atoms that are constituents of torsion angles, (a,, by, c, d)) with i e and determine the current torsion angles using T. Given any atomic configuration
Drot, using SetSubst wuctttutttttttttotto X created using this method, we can reach Xo by iteratively transforming X using R(-T(a,, by, C, d4) ) for every i, Iterative application of this transformation is typical (Jing et al., 2022) in other work and will bring to i 0 for all torsion angles. This is the starting point for Torsional Adjoint Sampling! By setting the torsion angles to zero, but retaining the atom types, molecular topology, bond lengths, and bond angles, we have an initial sample to transform with stochastic control.
Our SDE evolves in DD starting at Po = 0 and ending at 1 with the standard behavior of an SDE on a flat torus. However, our drift estimator of the torque is a function of Xx, implying we need to update the X, positions as yor evolves under the SDE.
25

---

Adjoint Sampling: Hiighy--Sallabl Diffusion Samplers via Adjoint Matching
Given drift estimate t(x,) from our model, update positions X1+1 D., R'x (t. ea) with e, (0, a (1) At the increment between steps and (t) is determined by the geometric noise schedule in (44) as a function of time.
Regression target In Cartesian Adjoint sampling, the regression target is V where the gradient is taken in Cartesian space. For our problem, we need to compute regression target V X1). These two quantities are related by the Jacobian 2x namely Vxg((Xi) of the function X = X1 applied to the V (X1) vector determined by energy model. To clarify, when we clip the gradient for the torsion model, we clip (X1)).
de (X1)) We compute this vector-Jacobin product using automatic differentiation ax
Torsion space posterior Adjoint sampling's innovations include fast sampling of points along the control trajectory using posterior. In the torsion case, we sample (rr from the posterior p (5p (01). This is done by affine transformation of the coordinates (0, 1] of the flat torus distributions from Section B.2 to x).
Transformation of torque! { When send X + -X, the torque is transformed t - -t, such quantities are called pseudovectors and their signed magnitudes are called pseudoscalars. Following Jing et al. (2022) and the implementation in Geiger et al. (2022), our networks predict dnou pseudoscalars without data augmentation.
G.3. Coverage Recall and Precision Metrics
Root Mean Square Deviation (RMSD) Similarly to Ganea et al. (2021) and Jing et al. (2022), we measure the so-called Average Minimum RMSD (AMR) and Coverage (COV) for Precision (P) and Recall (R). When measuring these metrics, we generate twice as many conformers as provided by CREST reference conformers. For K 2L let {Ct )En,, and {C& ]kj]I.K] be the sets of ground truth and generated conformers respectively. In our evaluations, we use = max L', 128), where T'' is the number of reference conformers given by CREST, taking the lowest energy conformers as a subset. The RMSD metric finds the best average distance between atoms of molecule with respect the reference molecule, taking into account all possible symmetries.
Coverage Recall
1
COV-R(0)
11 (1,.......... ......................
........,:: L. 3K (1, (I,..., K}, RMSD CT) 831
(84)
L
AMR-R = L
Σ min
RMSD(CR. CT)
(85)
kE(L....Kk
IE IE(I......))
Coverage Precision
COV-P(5) = [k E (1, (1....,K))
K): 31 L], RMSD(C CJ) < <8H
(86)
AMR-P = Σ
min
RMSD(C&, Cf)
(87)
kE kE(I....KK) IE(1.....)
where 8 > 0 is the coverage threshold.
G.4. Hyperparameters and Architecture Details for SPICE and GEOM-DRUGS
Cartesian Adjoint Sampling: We use an Equivariant Graph Neural Network (EGNN, Satorras et al. 2021) with 12 layers and a hidden feature dimension of 128. The model is trained for 5000 outer-loop iterations, sampling 100 batches each iteration from the replay buffer. Each GPU maintains own buffer with max size of 64000 samples and each iterations we generate 128 8 new molecules and energy evaluations for the buffer across 8 GPUs. The model is trained with batch size of 64 per GPU and follows geometric noise schedule with min 10-3 and o max = 1. Weight gradient clipping of 20 is applied. Additionally we ues the temperature 7 5 x 10-3 and regularization constant energy function. The regression target, the temperature-scaled gradient of the energy function is clipped at l2 norm 150. We use a step size of dt = 0.001 (i.e. 1000 nfe) for Euler-Maruyama SDE integration during training and evaluation.
100.0 for the regularized
Torsional Adjoint Sampling: We use a different equivariant graph neural networked called e e3nn (Geiger & Smidt, 2022; Geiger et al., 2022) with 6 convolutional layers with scalar features and 8 vector features each layer. There
26

---

Adjoint Sampling: Heighly-Sallable Diffusion Samplers via Adjoint Matching
is additionally a specialized layed called a pseudotorque layer by Jing al. (2022) that we use output to predict the pseudoscalar torque quantities. Our edge and node embeddings have 32 scalar features. The cutoff radius is 5.0 and we construct 50 Gaussian radial basis functions to embed distances between atoms, We use 8 GPUs with 1024 x 8 initial molecules and each iteration we generate 128 X 8 new molecules and evaluate their energy (gradient), which are then added to the GPU buffers. The model is trained for 3000 outer-loop iterations, sampling 100 batches each iteration from the replay buffer with max size 64000. The model is trained with a batch size of 64 and follows a geometric noise schedule with min = π x 10-2, max = n. and a "number of tiles" truncation parameter of 6 to approximate y on the flat tori (Appendix G.2). Weight gradient clipping of 10a0 is applied. Additionally we use the temperature T = 5 10 and the direct energy function without any regularization. The regression target, the temperature-scake gradient of the energy function w.r.t. torsions is clipped at e norm 200. We use a step size of dt = 0.01 (i.e. 100 nfe) for Euler-Maruyama SDE integration during training and evaluation.
H. Data Preparation for the Conformation Benchmark
Adjoint Sampling, in this context, takes in SMILES strings and outputs conformations. Here we explain the splits of the input SMILES strings and representative coordinates samples of conformations. Additionally, we explain the processes we used to produce, relax, and deduplicate our reference atomic conformations with RDKit (Landrum, 2013), CREST (Pracht et al., 2024), and ORCA (Neese, 2012) by following an extremely similar procedure to that of GEOM-DRUGS (Axelrod & Bombarlom 2022).
We are releasing the data we used in this paper as a challenging Conformation Benchmark with the goal of fostering development of scalable, amortized sampling algorithms. The data has three parts:
A training split of 24,477 SMILES strings from SPICE (Eastman et al., 2023). These define the molecular topology for the target molecule, but do not have any atomic coordinate information. The molecules have differing levels of flexibility with between O and 18 rotatable bonds?.
A test split' of 80 SMILES strings from SPICE divided into groups of 10, each with between 3 and 10 rotatable bonds. We also release 44.448 DFT annotated and geometry optimized conformations.
A test split of 80 SMILES strings from GEOM-DRUGS divided into groups of 10, each with between 3 and and 10 rotatable bonds. We identify 7024 conformations corresponding to those SMILES strings.
The details of are described below in Appendix H.1. Since we are repackaging data from SPICE and GEOM-DRUGS, our main contributions include (a) the organization of the benchmark and (b) the computational effort of finding 44,448 conformations from our SPICE test split using a similar procedure to that of GEOM-DRUGS.
One utilizes the benchmark by applying sampling algorithm to find conformers of given molecular topology and determining whether the algorithm recovers the samples we provided. In Figure 2, we report this as Coverage Recall % versus RMSD. We additionally provide tabulated values in Table 2. The conformers we provide should be treated as held-out data generated using an effective and expensive standard method. We do not claim to have comprehensively sampled the conformation space for these molecules, but we argue that they are a representative sample given modern methods. Furthermore, since we only release local minima, the sample does not come from Boltzmann distribution at a fixed temperature, i.e. the samples are not drawn from a probability distribution. There is no sense of a conformer being more "likely" than another in probabilistic manner in our data. Instead, the benchmark represents loosely defined notion of sampling by finding many low energy conformations in attempt to cover the configuration space.
H.1. SPICE Dataset
We used the PubChem subset of the SPICE dataset (Eastman et al., 2023) train and evaluate our models. This subset includes small, drug-like molecules comprised of between 18 and 50 atoms with elements Br, C, CI, H, N, O, P. and S. We are interested in conformations of molecules, implying that we care about the rotational degrees of freedom. We identify two different sets describing these degrees of freedom: number of rotatable bonds and torsional degrees of freedom, or
Rotatahle bonds are defined by RDKi t's function Cal cNumootatabl BBond (I (Landrum, 2013).
'The test split is disjoint from the train split and their union is all SPICE.
27

---

Adjoint Sampling: Hlighy-Seallbble Diffusion Samplers via Adjoint Matching
1089
4000
3435
2354
2045
3000
5000
2708
2130
3124
20000
2000
1300
A27
300
265
100
30C
10
12
13
35
36
37
38
0
24
15
19
37
19
19
1 cwwmRotatableeonns
ructMathhss with SMARTS
Figure 4. Histograms of rotatable bonds and torsional degrees of freedom in the SPICE training set. Left: Number of rotatable honds determined using RDKit's Ca hatuuumkotttableeenndna Right: Number of torsional degrees of freedom determined by Get ructMat ches with SMARTS string !!$ (+#*) &DD -218 !
torsions for short. In our case, rotatable bonds C torsional degrees of freedom. We compute the number of rotatable bonds using the RDKit package CalcWumootatableeBBeee
We count the number of torsional degrees of freedom using the SMARTS string [!$ + #+) !DD]] ! !!$ $ +) & [DD]] and Get chec,,, also from RDKit. We group molecules into classes according to the number of rotatable bonds, but we provide the torsional degrees of freedom to our Adjoint Sampling Torsion model. The distributions of rotatable bonds and the comparison between the two sets of rotational degrees of freedom can be found in Figure and Figure 5. We find fewer rotatable bonds than torsions matching substructure query, which is expected. More information about torsional model is in Appendix G.
We randomly subsampled 10 molecules from each class of molecules with a fixed number of rotatable bonds, in the range of 3-10, to create the test set. We performed pre-processing and CREST (Pracht et al., 2024) on the resulting 80 molccules. We also removed 30 additional molccules from the training set; 10 from each of the classes where the number of rotatable bonds was between 0-2. These do not appear in the test set since they are quite simple. Our procedure follows GEOM-DRUGS (Axelrod & Gomez- Bameaz-om,arrlli, 2022).
During pre-processing, first used RDKit generate 50 conformers with EmbedMul www.edMMlipllleooooocc using parameters pruneRma Thresh=0 01, maxAttempt s-5, useRandomCoords Corens-aodd-aallooo corresponding to pruning threshold of similar conformers of 0.01 Å, a maximum of five embedding attempts per conformer, coordinate initialization from the eigenvalues the distance matrix, and a random seed. If no conformers were successfully generated then numCon f.s was increased to 500, Afterwards, the conformers were optimized in the RDKit default MMFF force field. We deduplicated the conformers using RDKit's Bettess unless the calculation took longer than 48 hours then we switched RDKit's AlignMol. (97 used GetBestRMS, used Al ing..))). After removing the duplicate conformers that exhibited an RMSD 0.1A,, the ten conformers with the lowest energy were further optimized with approximate energy model known as extended tight binding (xTB) (Grimme et al., 2017: Friede et al., 2024). The conformer with the lowest xTB energy was used as input to the CREST simulation. We used the default hyperparameters from CREST version 3.0.2 including a 6 kcal/mol cutoff' on final conformers. Any SMILES strings containing slash or backslash indicate cis/trans stereochemistry and were skipped since their SMILES string is not preserved during the above procedure.
The output of the above procedure was thoroughly sampled potential energy surface by the standards of modern methods; however, the structures are optimized using the xTB level of accuracy. We took the process a step further and performed geometry optimizations using DFT with ORCA on each conformer in the test set. Specifically, we took the output of CREST and ran ORCA relaxation with settings:
Opt WB97M-V def2-tzvpd RIJCOSX def2/J NoUseSym DIIS NOSOSCF Norma Coon DEFGRID2 NormalPrint AllPop
8sef Convergence Tight maxiter 300 end selprop Dipole true Quadrupole true end Soutput: Print [P Reduceddrrboopoo_.] 1 Print [P P___eecceddrbboopo.M. 1
28

---

Adjoint Sampling: Seggll--Sallable Diffusion Samplers via Adjoint Matching
Print [P_ 1 Print P (P_BondOrder__M) Print [P_Fockian] 1 Print [P_orbEn] 2 end
&geom MaxTter 100 end
8scf THRESH 1e-12 TCUT 1e-13 end
Some lines were broken in the latex version of the code snippet above due to space limitations. These results are run at a very high level of DFT theory. Further geometry optimization may be desired at DEF DEFGRID3 with Tit however, such an extreme precision was prohibitive given our resources available at the time.
After the geometry optimization, we utilized RDKit's GetALIConforrerrBe pairwise duplicates and remove them from the dataset. The average number of reference conformers for each number of rotatable bonds is presented in the left side of Figure 6.
RMS and a threshold 0.1 À identify
We investigated the reason that the 80 molecule test set from SPICE generated an order of magnitude more conformers than the 80 molecule GEOM-DRUGS test set. We found, after replicating the GEOM-DRUGS conformers, that the molecules in SPICE test set are inherently more flexible than the ones from the GEOM-DRUGS test set. GEOM-DRUGS molecules tended have more rings and hydrogen bonds, leading to fewer conformations. In other words, SPICE tended to have flatter potentially energy surface. Since we followed the GEOM-DRUGS procedure, we applied the same kcal/mol cutoff' in CREST and kept many more conformers from SPICE. For Coverage Recall %, capped the reference conformers molecule at 512.
0.8
TerSonal 13 Marchee 53
10
27
16
0.6
76
250
419
343
0.4
A01
S05
10
61
1754
392
2595
0
156
2138
0
0
5 Number rotatable boncs using Calctumfctttsseeeonne
6
10
11
12
13
14
16
17
18
15
0.0
Figure Correlation heatmap hetween the torsional degrees freedom (y-axis) versus the number of rotatable honds (x-axis) in the combined SPICE train and test sets. Torsional degrees of freedom identified using GetSubs rubstrtttttttttt ches while rotatable bonds are identified using Colors indicate conditional probability: for each number of rotatable bonds (x-axis), the color shows the fraction of molecules with given number of torsional degrees of freedom (y-axis). Numbers in cells show the raw counts.
ableBBnda There are always more or equal torsional degrees of freedom than rotatable bonds.
H.2. GEOM-DRUGS Dataset
We also evaluated our models on the GEOMDRUGS dataset (Axelrod & Gomez-Bombarell. 2022). We took the test set from Torsional Diffusson (Jing et al., 2022), and excluded molecules that have elements other than Br, C, CI, F, H, N, O, P. and S. From the remaining molecules, we randomly sampled 10 molecules for each number of rotatable bonds in the range of 3-10 (calculated by RDKit's Ca attumoottaaalleBodda. The smallest molecule was comprised of 19 atoms and the largest of 65, slightly larger than SPICE.
29

---

Adjoint Sampling: Heibll-Scallabl Diffusion Samplers via Adjoint Matching
Iwernge Marder Centermee
129
100
1011
225
42.2
240.
11.9
Numbor Rotzzaicc Bonda INDRI
SPICE
GEOM-DRUGS
Figure 6. Average number of conformers in the SPICE (left) and GEOM- DRUGS (right) test sets vs number of rotatable bonds calculated by RDKit's function Cal cNumRot ataalatabnaaatoea Note the different y-axes between the figures.
0.00
0.54
0.00
Number
Number Aums
SPICE: number of atoms.
SPICE: number of heavy atoms (excludes Hydrogen).
Figure Atom counts appearing SPICE (train test) dataset: (left): Frequency total atoms. (right): Frequency of only heavy atoms (excluding Hydrogen). Our conformer generation predicts positions of all atoms.
We the conformers in the GEOM-DRUGS dataset the reference for evaluation. The conformers represented as RDKit molecule object. To use the same format the SPICE dataset, extracted the Cartesian coordinates from the RDKit molecules and wrote them into xyz files for each molecule. The average number of reference conformers for each number of rotatable bonds are presented in the right side of Figure 6.
30

---

Adjoint Sampling: Highly-SSalabble Diffusion Samplers via Adjoint Matching
I. Additional Experimental Results
I.I. Energy Histograms for Synthetic Energy Experiments
In this section, we provide additional qualitative comparisons between the energy distributions generated by our synthetic energy benchmarks and those obtained via ground -arunh MCMC sampling. Appendix I.1 illustrate histograms for three distinct systems DW4, LJ13, and LJ55. These results clearly indicate that Adjoint Sampling effective avoiding high-energy regions, especially in the case of IJ55, which features the highest-dimensional energy surface.
MONCERR
DOX
0.30
BCX
1120
0.30
1.312
0.3
0.25
24
72
-14
0.31
60
-40
030
150
-30
-30
-00
-33
100
-50
(a) DW4
(b) LJ13
(c) LJ55
Figure Energy Histograms of the Synthetic energy benchmark overlayed the ground-truth MCMC generated samples.
1.2. Runtime Analysis of Molecular Conformer Generation
This section further examines the computational efficiency of our proposed Adjoint Sampling technique through a detailed runtime analysis of molecular conformer generation task.
As illustrated in Figure 9, Adjoint Sampling overcomes two significant computational bottlenecks found in prior approaches, namely the intensive SDE simulation and the costly energy evaluation steps. By mitigating these challenges, our method able perform substantially more gradient updates within the same execution time, leading enhanced overall performance in tasks with expensive energy function evaluation, such conformer generation.
SDE Simulation
Adjoint Sampling
Energy Evals
RAM iDEM Updates x100
SDE Backprop x1
iDEM w/ 8 MC samples
iDEM w/ 64 MC samples
Discrete Adjoint (PIS)
2
5
Runtime (%)
Figure 9. Conformer Generation Run-time: A run-time analysis breakdown of Adjoint Sampling compared seleeted basclines at the level of gradient updates. Adjoint Sampling overcomes the two primary bottle-knecks of previous methods: SDE simulation and energy evaluation, allowing Adjoint Sampling perform far more gradient updates in the run-time.
I.3. Choice of energy function
The eSEN energy model (Fu et al., 2025), used in this work, has been shown to accurately predict DFT energy and forces on the SPICE dataset. We use an eSEN network with 2 layers, Lmax = 2, .max 0. 64 channels, and radius cutoff of 4.5 X. This model has around 900k trainable parameters. To demonstrate the robustness of the adjoint matching algorithm to the choice of the energy function, we conduct additional molecular conformation sampling experiments using adjoint matching
31

---

Adjoint Sampling: Hiighy-Scallbbl Diffusion Samplers via Adjoint Matching
Table Test MAE for cSEN-900k and MACF-OFF-23-M the SPICE-MACE-OFF dataset. Energy (E) MAE in me meaomm
Force (F) MAE is in V/A.
PubChem
DES370K M.
DES370K D.
Dipeptides
Sol. AA
Water
QMugs
Model
E
F
F
E
E
F
F
E
F
cSEN-900k
0.66
14.37
0.50
5.32
0.52
5.62
0.37
8.52
1.41
17.16
0.71
12.66
0.44
14.56
MACE-OFF-23-M
0.91
20.57
0.63
9.36
0.58
9.02
0.52
14.27
1.21
23.26
0.76
15.27
0.69
23.58
on pretrained OFFE-- energy model (Batatia al., 2022; Kovács et al., 2023). The test-set prediction errors of eSEN-900k and the MACF-OFF-23-M (Kovács et al., 2023) model reported in Table 3. The same energy model is used for both relaxation and for optimizing adjoint sampling, where applicable. Table 4 shows that adjoint matching with MACE-OFF-M and eSEN achieve similar performance in generating molecular conformation.
Table Recall and precision metries for conformer generation using different energy models. Coverage values thresholds 1.25Á. Standard deviations are computed across molecules in the test set. (These experiments were done with an earlier version of the code.)
SPICE
Recall
Precision
Method
Cov. f
AMR
Cov. f
AMR
cSEN
80.28 22:8
0.96 2025
47.97
1,25 poor
MACF
74.93_mw
1.02_mm
50.22_mm
1o28
eSEN (tpretrain)
88. 19.0
0.85 =025
65.12:22.00
1.13
MACE (+pretrain)
89.62 19.25
0.84-0z
67.07 #2244
1.11 10.34
w/ relaxation
eSEN
91.49-18.2
0.71
60.00=25.7
1.02
MACE
91.06.mm:
0.72
65.19 (2622)
1.05
eSEN (+pretrain)
96.36 8.9
0.60
76.46
0.92
MACE (tpretrain)
95.73-1132
20
77.08
222.8
0.92=0.4
32

---

Adjoint Sampling: Highly-Scalabll Diffusion Samplers via Adjoint Matching
1.4. Ablating Reciprocal Projection
Here present the same ablation of the Reciprocal Projection as performed for the Synthetic experiments in Table but for the molecular conformer generation task. Recall that instead of using the Reciprocal projection (RP) to sample X, given X1, we simply store sample pairs (4) in the buffer and train the Adjoint Matching objective (12). We call this Adjoint Sampling w/o RP, which helps demonstrate the effectiveness of the Reciprocal projection. For the purpose of this ablation, we only consider the MACE energy and do not perform any post generation relaxation.
Table Recall and precision metrics for Adjoint Sampling with and without use the Reciprocal Projection (RP). Coverage values are for thresholds of 1.25A. Standard deviations are computed across molecules in the test set. We see that the Reciprocal Projection results in significant improvements across all metrics, especially when generalizing to the unseen dataset GEOM-DRUGS. No pretraining and no relaxation is performed post generation. (These experiments done with earlier version of the code.)
SPICE
GEOM-DRUGS
Recall
Precision
Recall
Precision
Method
Cov. t
AMR
Cov,
AMR $
Cov. f
AMR
Cov.
AMR 1
Adjoint Sampling w/o RP (Ablation)
78.25 1006
0.98 2034
57.64_12:0
1.23 20:4
53.4629677
2.30
36.16
1.74
Adjoint Sampling (Ours)
80.28
2.96
47.97
1.25
61.36
1.17
41.19
1.46
33

---

Adjoint Sampling: Hiihhl--callabl Diffusion Samplers via Adjoint Matching
I.5. Threshold Ablation for Molecular Conformation Generation
AdSempiig
pling 1+preirain)
AdSenpling
AcpSumpiing (+pretrain)
1.0
100
Coverage Recall (%)
Coverage Recall (%)
80
20
Number Rotatable Bonds
Number of Rotatable Bonds
SPICE (Threshold=0.75A)
SPICE (Threshold=1.25
RDKC
Tonional AdSSmmping
Cartesian AdfSSmmpiing [+pretrain)
RDK
Torsional AdSempling
Cartedden AdfSumpiing (+pretrair)
100
Coverage Recall (%6)
Coverage Recall (%6
60
60
56
Number Rotatable Bonds
Number of Rotatable Bonds
GEOM-DRUGS (Threshold=0.75A)
GEOM-DRUGS (Threshold=1.255)
Figure 10. Coverage recall (%) after relaxation for molecules of different number of rotatable bonds. see that the gap between our model and RDKit increases for higher number of rotatable bonds.
34


Reviews:
[{"review_id": "VviO5jZXpR", "reviewer": "Reviewer_51Yn", "summary": "This paper proposes a novel neural sampling method, Adjoint Sampling, based on stochastic optimal control (SOC) and recently published adjoint matching method. The proposed method uses reciprocal projections alternating with reciprocal adjoint matching, and allows for incorporating the key symmetries from the considered energies. Adjoint Sampling allows for a significant reducing the number of needed energies evaluations per gradient step. Moreover, the paper introduces novel benchmarks for conformer generations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We’ve incorporated new baselines and metrics, and have produced figures to better illustrate our claims.\n\nAdditional figures & results: https://sites.google.com/view/adjointsamplingrebuttal\n\n> I’m not sure if the scalability was properly based on convincing experimental setting\n\nWe agree we did not emphasize this enough. Please see the link above for a runtime plot. Adjoint Sampling is the only method so far that can work with $(X_t, X_1)$ pairs — no full trajectories — and can scale to computationally expensive energy functions. Further details are provided in the answers of the reviewer’s other questions.\n\n> I strongly encourage the authors to include NLL and ESS metrics for the experiments on LJ potentials and DW-4. LJ potentials experiments should be equipped with the sampled energies histograms compared with the ground truth energy histogram.\n\n> I would recommend adding other samplers like FAB [1], or DDS [2] as the baselines for experiments\n\nWe have added several new baselines, including DDS. We have also included path-based ESS and ELBO estimates, and energy histograms. New results and figures can be found in the link above. For ESS and ELBO, we used the same method as [1,2,3]. However, we were not able to get reasonable estimates for iDEM at this time following the approach of [1] so have decided not to include them for now.\n\n> I think that putting more attention on a recent approaches for scaling the training or abilities of diffusion samplers, would be beneficial for this work, e.g., [1] or [2].\n\nWe completely agree! We will add more discussions to existing methods.\n\nWe note that there may be a conflation of “off-policy” and “scalability”, which are best decoupled. Off-policy methods, such as the log-variance divergence [4] and trajectory balance [5], do not take gradient of the energy function (as they do not differentiate through the model sample) and actually strongly relies on parameterizing the drift using the gradient of the energy function as discussed in [6]. In our setting, the energy function is more expensive than the drift network (see runtime plot), so our experiments do not use the energy function as part of parameterizing the drift. Another major downside is that these methods require the full trajectory, only being able to use sub-trajectories if the time-marginals are either learned [5] or prescribed [7]. In contrast, Adjoint Sampling is an on-policy method (resulting in a more direct update to the current model), directly works with the energy gradient rather than the energy values, and only requires $(X_t, X_1)$ pairs. \n\nIn our additional experiments, we have added a log-variance baseline from [3] to the synthetic energy experiments which did not scale beyond DW4. This is because the LJ potentials can have extremely large values which an importance-sampling based method like log-variance does not handle well (again, without taking the gradient of the energy into the parameterization).\n\n> It’s unclear to me why lean adjoint matching and adjoint sampling actually produce good gradients, since we are not minimizing the KL divergence anymore with the lean adjoint matching?\n\nInstead of optimizing the KL, Adjoint Matching (AM) directly regresses onto the control in the manner of a consistency loss. AM is still a fixed point iteration method and has the optimal control as the unique solution, and can be interpreted as removing a stochastic term from the KL gradient that has expectation zero at the optimum (as discussed in the original AM paper). In particular, the lean adjoint state does not depend on the learned control and only the base process.\n\nThis last property of the lean adjoint state is incredibly important in the sampling setting, and in the paper we use simple base processes so the lean adjoint state has a closed-form solution. This leads to allowing us to use only $(X_t, X_1)$ samples for training, and the Reciprocal projection further enhances optimality. We plan to provide more details about these derivation steps in the main paper.\n\n[1] “Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling”\n\n[2] “Path Integral Sampler: a stochastic control approach for sampling”\n\n[3] “Improved sampling via learned diffusions”\n\n[4] “From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training.”\n\n[5] “Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization”\n\n[6] “No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers”\n\n[7] “Sequential Controlled Langevin Diffusions”", "claims_and_evidence": "This paper introduces Adjoint Sampling and claims that is more efficient, highly-scalable and theoretically grounded. I agree that most of the claims have good evidences. However, I’m not sure if the scalability was properly based on convincing experimental setting, and the lack of other baselines (e.g., iDEM) on conformer prediction experiments.", "methods_and_evaluation_criteria": "The proposed evaluation setup is mostly typical for neural samplers community, but the number of baseline methods is limited. Moreover, the paper introduces novel benchmarks on conformer predictions. However, I strongly encourage the authors to include NLL and ESS metrics for the experiments on LJ potentials and DW-4. LJ potentials experiments should be equipped with the sampled energies histograms compared with the ground truth energy histogram.", "theoretical_claims": "I agree with the authors that the method is interesting and theoretically grounded. As far as I’ve checked, I haven’t found any obvious flaws in the theoretical parts.", "experimental_designs_or_analyses": "Overall, I think that the experimental setting is reasonable. I would recommend adding other samplers like FAB [1], or DDS [2] as the baselines for experiments, and add the previously mentioned metrics.\n\n**References:**\n\n[1] Midgley, Laurence Illing, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. \"Flow Annealed Importance Sampling Bootstrap.\" In The Eleventh International Conference on Learning Representations.\n\n[2] Vargas, Francisco, Will Sussman Grathwohl, and Arnaud Doucet. \"Denoising Diffusion Samplers.\" In The Eleventh International Conference on Learning Representations.", "relation_to_broader_scientific_literature": "This paper is relevant to the neural samplers community and proposes a novel method, based on very recently presented Adjoint Matching approach. However, I think that putting more attention on a recent approaches for scaling the training or abilities of diffusion samplers, would be beneficial for this work, e.g., [1] or [2].\n\n**References:**\n\n[1] Berner, Julius, Lorenz Richter, Marcin Sendera, Jarrid Rector-Brooks, and Nikolay Malkin. \"From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training.\" arXiv preprint arXiv:2501.06148 (2025).\n\n[2] Sanokowski, Sebastian, Wilhelm Berghammer, Martin Ennemoser, Haoyu Peter Wang, Sepp Hochreiter, and Sebastian Lehner. \"Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics.\" arXiv preprint arXiv:2502.08696 (2025).", "essential_references_not_discussed": "Please, refer to the previous sections and already mentioned references.", "other_strengths_and_weaknesses": "**Strengths:**\n\n[1] Introducing the novel sampler method with the strong theoretical results and good empirical evidence of its properties.\n\n[2] Presenting a novel conformal generation benchmark.\n\n\n**Weaknesses:**\n\n[1] Limited evaluation in terms of low number of baselines and missing metrics.\n\n[2] Missing important references to other works.", "other_comments_or_suggestions": "For other comments, please refer to the previous sections.", "questions_for_authors": "**Questions:**\n\n[1] It’s unclear to me why lean adjoint matching and adjoint sampling actually produce good gradients, since we are not minimizing the KL divergence anymore with the lean adjoint matching?\n\nFor other questions, please refer to the previous sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wdIqycrnqE", "reviewer": "Reviewer_MRRW", "summary": "This paper introduces Adjoint Sampling, a novel framework for efficiently sampling from an unnormalized density function. The framework reformulates the sampling problem as a stochastic optimal control problem. Building on the adjoint matching method, the authors propose the Reciprocal Adjoint Matching method, which allows for multiple gradient updates without requiring evaluations of the energy model. Additionally, the authors provide theoretical results to validate the convergence of the Reciprocal Adjoint Matching method. Empirical evaluations on a molecular conformation sampling task demonstrate the effectiveness of the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for being candid and providing us the opportunity to substantiate our claims, which we believe we can do. The reviewer is concerned with our claims regarding (i) the efficiency and analysis of Adjoint Sampling and (ii) our proposed large-scale sampling benchmark. We agree that we did not emphasize these contributions enough, and hope to clarify them in the following answers.\n\nPlease find additional ablations, experiments, and runtime figures here: https://sites.google.com/view/adjointsamplingrebuttal \n\n> Analysis and comparison to Adjoint Matching\n\n**Efficiency from a simple base processes**\n\nFirstly, note that the original Adjoint Matching (AM) method is designed for general base process, and has the same computational cost as the path integral sampler, requiring full trajectory simulation and solving the lean adjoint state. One of our first key observations is that the lean adjoint state can be solved in closed-form when the base process is simple. \n\n**Ablation and analysis of the Reciprocal projection**\n\nFurthermore, we note that the Reciprocal projection further enhances optimality, allowing us to explore the search space much faster (see Figure 5 of above link) as it allows training on uncorrelated X_t samples. We performed the requested ablation where we removed Reciprocal projection (see “Adjoint Sampling w/o Reciprocal Projection” in Tables 1 & 2). We see that Adjoint Sampling **with Reciprocal projection** performs better in all settings. On the LJ classical force fields, it avoids high energy regions much better (better Energy W2 metric), and on our amortized conformer generation task it covers the target distribution much better (better recall and precision metrics). We hypothesize this is because the Reciprocal projection lets the model see far more diverse trajectories across different molecular conditionings. In the amortized benchmark, the number of samples per molecule is extremely small, so having a strong learning signal for each energy evaluation becomes very important. Moreover, we know that the Reciprocal projection is theoretically grounded and preserves the optimal solution as a unique fixed point. \n\n**Runtime plot**\n\nWe have included a run-time breakdown in Figure 2 of the link. It shows that Adjoint Sampling is computationally efficient in terms of time-spent evaluating the energy and sampling model. This is what enables us to scale up to both larger architectures and energy models, enabling us to sample from molecular energy foundation models. Other methods will either require full trajectory optimization (such as PIS, log-variance) or require an intractable number of energy evaluations (such as iDEM). \n\n> Large-scale experiments for sampling\n\nAs mentioned above, Adjoint Sampling is specifically designed for sampling from unnormalized distributions where a simple base process can be used. The text-to-image finetuning problem is a more general problem statement where a pre-trained generative model is used as the base process. Since our methodological contributions rely on the choice of a simple base process, we restrict ourselves to the pure sampling setting, where only an unnormalized density is given.\n\nIn the literature of sampling from unnormalized densities, our proposed benchmark is actually THE most difficult benchmark to date, requiring highly scalable algorithms. Prior related works have only experimented with classical (synthetic) force fields such as the Lennard-Jones experiments we have included. In contrast, the conformer generation benchmark uses a large graph neural network as the energy function, and hence each evaluation incurs a runtime cost. Furthermore, this benchmark is aimed at learning conditional sampling models, amortized over 24000+ molecules, with the test metrics being generalization to unseen molecules. This type of amortized benchmarks for sampling is incredibly rare and has not been well-explored. Due to this amortization, expensive methods that require full simulations for each gradient update (such as Adjoint Matching) does not scale well — again, note the runtime plot (Figure 2 in the link), Adjoint Matching in its basic form has a similar computation cost as the Discrete Adjoint (PIS).\n\nFinally, the reward fine-tuning benchmarks are aimed at finding one good image sample per text prompt. Here, our benchmark requires the model to find a set of diverse samples per molecule (ALL local minima), and the metrics of interest are precision and recall against a set of highly-diverse ground truth samples from density functional theory (DFT) that took weeks to obtain even on a large cluster. Performing well on this benchmark has direct consequences in advancing computational chemistry and drug discovery.\n\nWe hope this answers the reviewer regarding (i) why Adjoint Sampling works only for the sampling task, and (ii) how our proposed new benchmark is precisely encouraging the sampling community to look into harder and larger-scale problems.", "claims_and_evidence": "I believe the answer is no. While the authors claim that the adjoint sampling method is more scalable compared to the original adjoint matching method, a thorough analysis and comparison are necessary to evaluate the efficiency of the adjoint sampling method. Furthermore, experiments on large-scale datasets are needed to provide evidence of the method's scalability.", "methods_and_evaluation_criteria": "The method is well-suited to addressing the problem.", "theoretical_claims": "I have checked the proofs for the theorems.", "experimental_designs_or_analyses": "As I have mentioned before, experiments on large-scale datasets are needed.", "relation_to_broader_scientific_literature": "The method primarily combines the adjoint matching technique with the reciprocal projection approach.", "essential_references_not_discussed": "The related works are essential to understanding the context for key contributions of the paper.", "other_strengths_and_weaknesses": "**Strengths:**  \n   \n1. The writing is very clear and easy to follow.    \n2. The method is well-suited to the problem.  \n   \n**Weaknesses:**  \n   \n1. **Lack of large-scale experiments and efficiency analysis:** As mentioned earlier, a thorough analysis and comparison are necessary to assess the efficiency of the adjoint sampling method. The original adjoint matching method would serve as a reasonable baseline. Furthermore, experiments on large-scale datasets are needed to demonstrate the scalability of the proposed method.    \n     \n2. **Limited experiments on other domains:** The paper lacks experiments in domains beyond the one studied. For instance, including experiments on text-to-image generation tasks and comparing the results to the adjoint matching method could further highlight the paper’s contributions and broaden its impact.", "other_comments_or_suggestions": "Please refer to the sections above.", "questions_for_authors": "Please refer to the sections above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r5F1lu9zMb", "reviewer": "Reviewer_Nbi6", "summary": "The paper proposes an algorithm that uses  diffusion-models for sampling from unnormalized densities which is rooted in stochastic optimal control. The proposed method is based on the adjoint-state and the resulting form of the objective is particularly simple as it is requires a regression to the (scaled) gradient of the terminal cost of the SOC problem. Furthermore, the paper explains how various symmetries such as periodic boundary conditions are included. Lastly, the authors propose a novel benchmark for sampling molecular conformers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We understand and agree with the reviewer’s concerns regarding additional baselines and the lack of clarity around the proposed benchmark.\n\nTo respond, we (i) have additional experiments, (ii) expand the discussion to related off-policy methods, and (iii) clarify why the proposed benchmark is much harder than existing ones.\nSee: https://sites.google.com/view/adjointsamplingrebuttal\n> the log-variance loss has similar benefits as adjoint sampling\n\nThere are problems with the log-variance loss that inhibits scalability. Firstly, it requires full trajectories of length N, and N network evals, per gradient. Partial trajectories can only be used if the time-marginals are either learned [1] or prescribed [2]. In contrast, Adjoint Samping only uses pairs of (X_t, X_1) samples, and one network eval, for each gradient.\n\nSecondly, the log-variance loss doesn’t make use of the energy gradient information. For log-variance to work in high-dimensions, the gradient of the energy is used directly as part of the drift parameterization and it can fail without this [3], but we avoid this since our focus is on computationally expensive energy functions (such as the proposed benchmark). In our tests, the log-variance loss fails to learn on LJ because the potentials have extremely large values.\n\n> other sampling methods \n\nWe added multiple baselines, including log-variance, DDS, and ablations. Additionally, we have incorporated the suggested reference [4] to estimate ESS and ELBO metrics. Please see Tables 1 & 2 in the above link. \n> more established benchmarks\n\nWe believe the DW and LJ problems are very common and have been used by many prior works.We mainly focus on our proposed benchmark.\n> how high-dimensional are the proposed benchmark problems?\n\nWe agree this was not emphasized well and hope to remedy this.\nThe number of atoms range from 3-50 (median 39), sample dimension is 3x the number of atoms, while the conditioning dimension is quadratic (bonds). See Figs 3 and 4.\n\nThis is an amortized setting for finding the distribution of conformers for over 24,000 molecules (i.e., 24,000 conditional sampling problems). The benchmark is designed to test for **generalization to unseen molecules**, whereas most existing sampling benchmarks (such as DW, LJ, alanine dipeptide) only test performance on a single & cheap energy function.\n\nFurthermore, each molecule has a number of conformers (representing low energy regions), ranging from a handful to a few hundred (Figs 6,7 in submission). The benchmark is specifically designed to test if a sampling method can find ALL the conformers for every molecule. As such, the metrics of interest are precision and recall against a set of highly-diverse ground truth samples computed using density functional theory (DFT) that took days to obtain even on a large cluster. \n\nFinally, the energy function for this benchmark is a large graph neural network. As such, care must be taken in order to not incur computational cost by evaluating the energy function too many times. See Fig 2 for a runtime plot.\n\nTo our knowledge, this is the most difficult sampling benchmark so far, which we believe can aid in directing future research on sampling methods. Performing well on this benchmark has direct consequences in advancing computational chemistry and drug discovery.\n> other sampling method on the novel benchmark\n\nIt is VERY difficult to get reasonable performance on this benchmark. The only other method that we are aware of that can train with only pairs of (X_t, X_1) samples is iDEM, which we now include (see Table 2 of the link). iDEM is biased when few MC samples are used; it typically uses 512 MC samples (== energy evaluations) per gradient, which is prohibitive (see Figure 2). We also added an ablation for the Reciprocal projection (see Table 1 & 2 of the link); without it, the model performs worse, which we hypothesize is due to the lack of exploration (see Figure 5 in the link).\n> [...] Alanine Dipeptide, which is a common problem in the sampling community?\n\nWe did not specifically test the alanine dipeptide setup (which is 22 atoms and uses a classical energy function). The molecules in our benchmark are much larger, and our energy function is much more expensive as it is a GNN trained to approximate forces from DFT.\n> Does the proposed objective avoid mode-collapse such as e.g. relative entropy minimization?\nOur method is related to the reverse KL rather than the forward KL (i.e., relative entropy). Mode collapse is inherently difficult as there is no way to know where modes are without fully exploring the search space. Our method relies on the choice of base process to determine what region to search. Finally, we again note that the new benchmark is specifically designed to test mode coverage as it requires finding all minima.\n\n[1] https://arxiv.org/abs/2310.02679\n\n[2] https://arxiv.org/abs/2412.07081\n\n[3] https://arxiv.org/abs/2502.06685\n\n[4] https://arxiv.org/abs/2406.07423", "claims_and_evidence": "The authors claim:\n\n> It is the first of its kind in allowing significantly more gradient updates than the number of energy evaluations\n\nand \n\n> However, all of these methods\nare hindered by their computational requirements, including\nexpensive differentiation through the sampling procedure,\ncomputation of higher-order derivatives in constructing the\ntraining objectives, or the need for importance sampling\n(i.e. multiple energy evaluations).\n\nHowever, using the log-variance loss [1] has similar benefits as adjoint sampling. Can the authors comment on that?\n\n[1] Richter, Lorenz, and Julius Berner. \"Improved sampling via learned diffusions.\" arXiv preprint arXiv:2307.01198 (2023).", "methods_and_evaluation_criteria": "The reviewer cannot judge if the proposed benchmark is useful due to a lack of knowledge about conformer prediction. Moreover, the authors do not compare their method to any other sampling method which is confusing. It would be more convincing if the authors could show the performance of their method on other, more established benchmarks, see e.g. [1] for a recent study.\n\n[1] Blessing, Denis, et al. \"Beyond ELBOs: A large-scale evaluation of variational methods for sampling.\" arXiv preprint arXiv:2406.07423 (2024).", "theoretical_claims": "The proofs of the results were skimmed and, to the best of the reviewers knowledge, appear to be correct.", "experimental_designs_or_analyses": "The reviewer is familiar with the experimens apart from the new benchmark tasks.", "relation_to_broader_scientific_literature": "The paper proposes a novel objective rooted in SOC that has a simple form and allows for the usage of off-policy learning with a replay buffer. If the authors could demonstrate that their method consistently performs well on more established benchmarks, then the method could have a high impact.", "essential_references_not_discussed": "The most relvant references have been discussed to the best of the reviewers knowledge.", "other_strengths_and_weaknesses": "Weaknesses\n\n- The paper only consideres few sampling methods as baselines\n- The authors do not compare their method to any other sampling method on the novel benchmark\n\n\nStrenghts\n\n- The authors propose several geometric extensions, some of which are, to the best of the reviewers knowledge, novel in the context of sampling from unnormalized densities. \n- The authors propose a new benchmark for sampling problems (although the reviewer cannot judge wheter or not this benchmark will have an impace)\n- The simple regression objective suggest numerical stability and scalability", "other_comments_or_suggestions": "None", "questions_for_authors": "- Did the authors try the proposed method on Alanine Dipeptide, which is a common problem in the sampling community? \n- I might have missed it, but how high-dimensional are the proposed benchmark problems? \n- Does the proposed objective avoid mode-collapse such as e.g. relative entropy minimization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b5ryvVK4xI", "reviewer": "Reviewer_gL2K", "summary": "The paper proposes Adjoint Sampling, a scalable and effective method for diffusion samplers. Authors build their ideas on top of adjoint matching and propose several advancements for scalable and effective training of diffusion samplers. Experiment results validate that the proposed method outperforms several baselines across various benchmarks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for supporting our paper. Note that we have added additional experiments and figures to provide more insight into our work. See https://sites.google.com/view/adjointsamplingrebuttal. Below we answer the reviewer’s questions in detail.\n> It seems the training procedure can be conducted in an off-policy manner. \n\nWe wish to clarify that our method is not actually off-policy, in the sense that we cannot sample from an arbitrary SDE, which is what is often referred to as off-policy by other methods. That being said, there are many downsides to an off-policy method – such as requiring full trajectories for optimization, and our method is actually much more scalable due to requiring only $(X_t, X_1)$ samples.\n> Are there any other possibilities to improve the sample efficiency by using several off-policy training schemes?\n\nThe existing off-policy methods come with their own drawbacks, such as requiring full trajectories and not using the gradient of the energy for the training loss. We added additional experiments using the log-variance loss [1], but it does not perform well. We find that such off-policy methods only work well when the gradient of the energy is used directly as part of the drift parameterization [2], which is something we need to avoid as we move on to utilizing more computationally costly energy functions.\n> Is there any reason why the samples are uniformly sampled from the buffer? \n\nThis is a good suggestion. We did not specifically look into prioritized replay buffer, but it is possible to use importance weights to design a prioritized replay buffer where higher weights are sampled more frequently [3]. We note that such an approach is difficult to implement for our proposed (amortized) benchmark, which aims at solving 24,000+ sampling sub-problems using a single model, as the priority over samples needs to be computed for each sub-problem separately. In order to train on all sub-problems and be able to generalize, we cannot sample each sub-problem too many times as it incurs a high computational cost so we just used a simple uniform buffer.\n> It will also be nice to compare the time complexity or performance of naive adjoint sampling (e.g., without reciprocal adjoint matching).\n\nWe have included additional ablations without the Reciprocal projection (see Tables 1 & 2 of the link), and also a runtime comparison between different methods (Figure 2 of the link). We note that the naive Adjoint Matching in terms of computational cost is on-par with Discrete Adjoint (PIS) in terms of runtime, while the ablation “Adjoint Sampling w/o Reciprocal” is on-par with Adjoint Sampling in terms of runtime. However, without the Reciprocal projection, the performance deteriorates because it does not as efficiently search the sample space. With Reciprocal AM, the samples $X_t$ are uncorrelated, whereas without it, the samples across time are from the same trajectory (see Figure 5 of the link). We hypothesize that the improved performance when using Reciprocal AM (in terms of Energy W2 for Table 1 and Recall for Table 2) is due to the ability to see more diverse X_t samples during training.\n\n[1] “Improved sampling via learned diffusions”\n\n[2] “No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers”\n\n[3] “Sequential Controlled Langevin Diffusions”", "claims_and_evidence": "The claim made in the submission is supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The authors follow the conventional experiment setting for diffusion samplers.", "theoretical_claims": "I check that the theoretical claims of the paper (ex. Proposition 3.1) are correct.", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "Diffusion samplers can be applied various scientific applications such as physical simulations.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "There are several comments and suggestions listed below.\n\n- It seems the training procedure can be conducted in an off-policy manner. Is there any reason why the samples are uniformly sampled from the buffer? Are there any other possibilities to improve the sample efficiency by using several off-policy training schemes?\n\n[1] Sendera, Marcin, et al. \"Improved off-policy training of diffusion samplers.\" The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n- It will also be nice to compare the time complexity or performance of naive adjoint sampling (e.g., without reciprocal adjoint matching).", "other_comments_or_suggestions": "Please see above.", "questions_for_authors": "Please see above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}]

Author rebuttal:

Meta-review: The authors present AdjointMatching, an algorithm based on stochastic optimal control, to learn diffusion processes that sample from non-normalized probability densities, e.g. the Boltzmann distribution, falling into the broad category of neural samplers. Beyond the related work mentioned by the reviewers, the approach bears some conceptual resemblance to neural thermodynamic integration (Mate et al 2024, JPhysChemLett and prior work from these authors on interpolation of Boltzmann densities) while being motivated from a completely different perspective, and providing a more transparent regularization approach. The reviewers have a strong positive sentiment towards the paper, and following the discussion phase all doubts and questions initially raised by the reviewers particularly with regard to metrics reported, have been resolved.
Decision: Accept (poster)