Non-equilibrium Annealed Adjoint Sampler
Authors: Jaemoo Choi
Yongxin Chen
Molei Tao
Guan-Horng Liu
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion sampler
stochastic optimal control
annealed sampling
adjoint matching
Boltzmann density
Primary area: probabilistic_methods
TL;DR: We introduce the Non-equilibrium Annealed Adjoint Sampler (NAAS), a novel SOC-based diffusion sampler that leverages annealed reference dynamics

Abstract: Recently, there has been significant progress in learning-based diffusion samplers, which aim to sample from a given unnormalized density. Many of these approaches formulate the sampling task as a stochastic optimal control (SOC) problem using a canonical uninformative reference process, which limits their ability to efficiently guide trajectories toward the target distribution. In this work, we propose the **Non-Equilibrium Annealed Adjoint Sampler (NAAS)**, a novel SOC-based diffusion framework that employs annealed reference dynamics as a non-stationary base SDE. This annealing structure provides a natural progression toward the target distribution and generates informative reference trajectories, thereby enhancing the stability and efficiency of learning the control. Owing to our SOC formulation, our framework can incorporate a variety of SOC solvers, thereby offering high flexibility in algorithmic design. As one instantiation, we employ a lean adjoint system inspired by adjoint matching, enabling efficient and scalable training. We demonstrate the effectiveness of NAAS across a range of tasks, including sampling from classical energy landscapes and molecular Boltzmann distributions.

Non-equilibrium Annealed Adjoint Sampler
Jaemoo Choil.", Yongxin Chen Molei Tao', Guan-Horng L
GGoorgia Institute of Technology, "FAIR at Meta, "Core contributors [jcho1843, yongchen, mtao) @gatech. edu, ghliudmeta. com
Abstract
Recently, there has been significant progress in learning-based diffusion samplers, which aim to sample from a given unnormalized density. Many of these approaches formulate the sampling task stochastic optimal control (SOC) problem using a canonical uninformative reference process, which limits their ability to efficiently guide trajectories toward the target distribution. In this work, we propose the Non- Equilibrium Annealed Adjoint Sampler (NAAS), novel SOC-based diffusion framework that employs annealed reference dynamics as statioattonrry base SDE. This annealing structure provides a natural progression toward the target distribution and generates informative reference trajectories, thereby enhancing the stability and efficiency of learning the control. Owing to our SOC formulation, our framework can incorporate a variety of SOC solvers, thereby offering high flexibility in algorithmic design. As one instantiation, system inspired by adjoint matching, enabling efficient and scalable training. We demonstrate the effectiveness of NAAS across a range of tasks, including sampling from classical energy landscapes and molecular Boltzmann distributions.
employ lean adjoint
Introduction
fundamental task in probabilistic inference is draw samples from a target distribution v(x) given only its unnormalized density function, where the normalizing constant is intractable. This challenge arises across a wide range of scientific domains, including Bayesian inference (Box and Tiao, 2011), statistical physics (Binder et al., 1992), proteins (Jumper et al., 2021; Bose et al., 2024), and molecular chemistry (Weininger, 1988; Tuckerman, 2023; Midgley et al., 2023). Traditional approaches predominantly leverage Markov Chain Monte Carlo (MCMC) algorithms, which construct a Markov chain whose stationary distribution the target distribution N (Metropolis et al., 1953; Neal, 2001; Del Moral et al., 2006). However, these methods often suffer from slow mixing times, leading to high computational cost and limited scalability.
To overcome this limitation, recent research has focused on Diffusson Neural Samplers- fam- ily of Icarned stochastic differential equation (SDE)-based samplers in which the dynamics are parameterized by neural networks (Zhang and Chen. 2022; Vargas et al., 2023; Chen et al., 2024; Albergo and Vanden Eijnden,, 2024; Phillips et al., 2024; Chen et al., 2024). Among diffusion neural samplers, many of the prior SOC-based approaches (Zhang and Chen, 2022; Vargas et al., 2023; Domingo-Ena et al., 2025; Behjoo and Chertkov, 2025) aim to solve the stochastic optimal control (SOC) problem to construct unbiased diffusion samplers. However, these methods typically learn the controlled dynamics starting from a stationary or uninformative base reference, and therefore do not fully leverage annealed reference dynamics that progressively guide trajectories toward the target distribution.
Motivated by this limitation, propose Non-Equilirrium Annealed Adjoint Sampler (NAAS), new SOC-based diffusion framework that employs annealed reference dynamics as the base SDE. The annealing structure provides natural progression toward the target distribution and generates
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Table 1: Compared prior diffusion samplers, our Non-equilibrium Annealed Adjoint Sampler
(NAAS) effectively guides the sampling process toward the target distribution by leveraging annealed
reference dynamics, without using importance weighted sampling (IWS) during training.
Method
Require IWS
Annealed Reference SDE
PIS (Zhang and Chen, 2022) DDS (Vargas al., 2023)
No
x
LV-PIS LV-DDS (Richter Berner, 2024)
No
iDEM [Akhound-Sadeg) et al., 2024)
yes
AS (Havens et al., 2025)
No
X
PDDS (Phillips ct al., 2024)
Yes
CRAFT (Matthews al., 2022) SCLD (Chen et al., 2024)
Yes
NAAS (Ours)
No
informative reference trajectories, which substantially improve the stability and efficiency of learning the control. As demonstrated empirically, this design leads to enhanced sample quality and target matching, especially when compared with methods relying on static references. To the best of our knowledge, NAAS is the first diffusion sampler that integrates the stochastic optimal control (SOC) perspective with reference annealed SDEs.
Because our method is formulated within the SOC framework, can naturally integrate diverse solver choices (Nüsken and Richter, 2021; Domingo-Earich et al., 2025), allowing flexible adaptation to different optimization settings. In contrast to existing annealed samplers (Albergo and Vanden- Eijnden, 2024; Phillips al., 2024; Chen et al., 2024), which rely emportance-weightee sampling (IWS), NAAS supports solver choices that are independent of resampling- estimators. Among these possible instantiations, we adopt Adjoint Matching (Domingo-Enrich al., 2025)-a recently proposed, scalable, and practical solver that achieves low-variance gradient estimation and stable optimization in digeesinnnii SOC problems.
As shown in Figure 1, our SOC formulation can be decomposed into two complementary SOC subproblems: (i) learning prior distribution J from the initial distribution &n with standard base reference, and (ii) transporting to the target v using controlled annealed dynamics. Together, these components define a controlled diffusion process that realizes the transport 6 μ - v, resulting in a sampler whose terminal distribution matches the target. In practice, we employ an alternating scheme update the two control functions associated these subproblems. We evaluate the effectiveness of NAAS on standard synthetic benchmarks and challenging molecular generation task involving alanine dipeptide, demonstrating stable convergence and strong sample quality across settings.
Our main contributions are summarized as follows:
We introduce SOC-based unbiased sampler where the reference dynamics is governed by an annealed SDE. This formulation allows sampling to begin from a well-behaved initial process that gradually transitions toward the target distribution.
We develop an efficient optimization algorithm based adjoint matching, which provides sample-efficient, and scalable approach to solving nigh-dimensiooal SOC problems.
We assess our method on standard synthetic benchmarks and an alanine dipeptide molecular generation task. Our approach consistently outperforms existing baselines, demonstrating superior sample quality and diversity.
2 Preliminary
Throughout the paper let be the target distribution in space X, and let U1 X + R be an energy function of the distribution v, i.e. v(-) e-U1()
Stochastic Optimal Control (SOC) The SOC problem studies an optimization problem subjected to the controlled SDE while minimizing a certain objective. We focus on one type of SOC problems formulated as (Kappen, 2005; Todorov, 2007)

---

Armeating Dyn.
VG(X,edt adlW
Nop-equi Bibbiemm Anneting Adjoint Sampler (NAAS)
Learn to sample from source
Learn samole torget
non
dX, mnp odW;
dX,
(2.4)
Leannd Controlled Dyn.
e10
Figure Visualzation of Training Process of NAAS. Given an anncaling reference dynamics that provide high-quality->99 yet imperfect-initias samples in high-density regions of the target distribution v, NAAS learns to sample from both the target v and the debiased source by alternate optimization between two control functions, u and va. with Adjoint Matching (eq. (21)) and Reciprocal Adjoint Matching (eq. (24)). This iterative procedure progressively aligns the sampling path with the optimal control plan, leading unbiased and efficient sampling from the target v.
min EE J.' (x.)) (((x.)) di + o(xi)
s.t. dx, = (X,) + or (x,) d/ + a,dW, Xa~1
(2)
where μ is initial distribution of random variable Xo. (x(),) is given reference drift and o; is the diffusion term. Here, we denote f: [0, 1] x X - R as the running state cost and X + IR as the terminal cost. Throughout the paper, we denote pase and p" as path measures induced by reference dynamics and optimal dynamics, respectively. The optimal control u" is given by
uo (x) = - o, (v.), where the value function V [0, x X + is defined as follows:
(1l f((x)))) di + 9(Xi) X, =
Vi(x) inf IE
s.t.
eq. (2).
Adjoint Matching (AM) Naively backpropagating through the SOC objective (1) induces pro- hibitive computational costs. Instead. Domingo-Enten Enrich et al. (2025) proposed scalable and efficient approach, called Adjoint Matching (AM), to solving the SOC problem in (1). The AM objective is built on the insight that the optimal control policy u* can be expressed in terms the adjoint process. Specifically, AM constructs a surrogate objective by aligning the control function u and quantity derived from the adjoint process. The precise formulation as follows:
min Exapa L/ [lux(X)) o,ā(1; di
stopgrad(a), stopgrad(u),
d
where dt at;
(a(t; x) + V(()).). a(1; X) = (x())).
Here, X ~ pd denotes full trajectories sampled from the controlled dynamics using the current detached policy ū, and â represents detached lean adjoint process a(t: X), which is obtained by backward dynamics starting from the terminal condition (g(Xi).
Reciprocal Adjoint Matching Recently, Adjoint Sampling (AS) (Havens al., 2025) propose more efficient variant of adjoint matching tailored for the special case where both the drift term

---

and the running cost vanish, i.e. b = 0 and f == O. In this setting, the lean adjoint a(t: X) admits closed-form expression; a(t: X) (XX)), which is constant over time. Furthermore, they observe that since the optimal controlled path measure satisfies the reciprocal property:
p p*(x) - ppp (Xo. X1)p base Xo, X1):
the conditional law intermediate states given the endpoints remains the same under both the optimal and the reference dynamics (Léonard et al., 2014). This insight simplifies the AM objective (4) to bypass the need to store full state-adjoint pairs {(x,, Itju... Instead, one can sample an base (-) Xo, X1) given only the boundary states intermediate state X, from the reference distribution P 10,11
(Xo, X1). Note that since the drift term vanished, i.e. b = 0, this conditional sampling is tractable.
SOC-based (Memoryless) Diffusion Samplers We now consider the SOC problem (1) under the specific setting where
plase (I) hase Po.1
/ / 0, g(x) log
(Xo, X1)
hase (Xo)pl hase (X1).
v(x)
The final condition implies that the initial and terminal states of the reference process are indepen- dent- property we refer to as memorylessness. Under this assumption, the value function at the initial time, Vz(x), becomes constant by Feynman-KK Kac formula (see Domingo-Enrich et al. (2025)). Then, applying Madoo-Nikodym derivative on path measures p" and p yields the following (Dai Pra, 1991; Pavon and Wakolbinger, 1991; Chen et al., 2021)
dpā,1 (Xo, (1),Xi dybar base (Xo, X1)
x x(x.)))) = (X2) [ c (=((xi)-V(XXo)
Po,1 (Xo, X1 )1xx
p e base (X1) v(X1).
That is, the terminal distribution induced by the optimal controlled SDE exactly matches the desired target Thus, solving the corresponding SOC problem leads to unbiased diffusion-based sampler. Notably, several recent diffusion samplers can be understood within this memoryless SOC framework:
Path Integral Sampler (PIS; Zhang and Chen, 2022) set μl = 6g. where the reference dynamics automatically becomes memoryless. The reference dynamics is defined as follows:
d.X, odW, 2. p (x) (x;
(8)
Several methods have been proposed to solve the associated SOC problem, including the adjoint system approach (Zhang and Chen, 2022), -aarianne loss (LV) (Richter and Berner, 2024), and reciprocal adjoint matching (Havens et al., 2025).
Denoising Diffusion Sampler (DDS; Vargas et al., 2023) This method constructs an approximately memoryless reference process using Ornstenn-Ullene is set as μ = (0, o² and the dynamics is given by:
(OU) dynamic. The initial distribution
Be x,dt + to BVB,,, ())) dx, = 2
(9)
where the annealing schedule 3e is chosen so that at Ja Bidt ss 0, ensuring approximate indepen- dence between Xo and XI.
3 Non-Equilibriun Annealed Adjoint Sampler
In this section, we introduce NAAS, novel SOC-based diffusion sampler that generates samples using non-equilibrium annealed processes. Specifically, NAAS is based on newly designed SOC problem featuring annealed reference dynamics (Section 3).. Since the problem requires sampling from an intractable prior μr, we present complementary SOC problem in Section 3.2 to approximate samples X, ~ j.. Finally, Section 3.3 provides a detailed description of the algorithm and its practical implementation details. A detailed proof for theorems in this section provided in Appendix A.

---

3.1 Stochastic Optimal Control with Annealing Paths
Let Uy and U1 be the energy potentials of some tractable prior (e.g.. Gaussian) and the desired target distribution v, respectively. We define a smooth interpolation between these potentials as follows:
U²(z)= = [tt, (z)),, with (z) = (0()), (==1(1) = (I),, (10)
where [tt.,,) denotes an interpolation cchmm--such linear or geometric interpolation. Based on this interpolation, we define an annealed SDE that serves as our reference dynamics:
dXx, = of (U,(X,)dI + o,dW,, Xo μl-
(11)
We denote a path measure induced by the reference process as plase, The following theorem introduces newly designed SOC problem whose solution yields sampler that transforms the specific prior distribution j to the desired target distribution
e-V: (-);
Theorem 3.1 (SOC-based Non-equilibrium Sampler). Let [Usheeoo1 be a fime-dependent potential function defined in eq. (10). Consider the following SOC problem:
min
E 1 + dt,
Ju
(12)
o 2
s.t. dx,
VUI(X))
to. dt + o,dWe Xo Xo~H μ4
where the initial distribution y is defined as p(x) x exp (xp(-UU(x)- Vo(=)),
(13)
and Vo(x)) is the value function associated with the corresponding path measure. Note that V V1(z) = 0 due to the absence of terminal cost. Then, the marginal distribution př (x) of the optimally controlled
process admits the following form:
(x) x - for all t E [0,
(14)
In particular, the terminal distribution satisfies pi(x) x exp (-U)(x)) =: v(x). Hence, the optimal controlled dynamics generates the desired target distribution v.
Sketch of Proof We outline the proof to provide an intuition behind the design of the objective function the SOC formulation of eq. (12). Suppose (x(x) e- Vvv and (x Pī (z)e (a(),
Then, the pair (v, Φ) satisfies the following PDEs, namely Hopf-Cole transform (Hopf, 1950):
(x) o² VU, V + Apt 0,Uipe po(z))o(r)= = p((), (x))(())
2
(15)
V (U + V U,
s (1) = v())
Note that the 0,U, terms in both PDEs arise from the running state cost in the SOC objective. The inclusion of this term allows us derive closed-form expression of ':
(x) Ce (U)(),
up some multiplicative constant C. By the definition, we obtain
pī (z) (x)( x e (x(())p(()) = exp(-U: V((z)),
which recovers eq. (14) claimed in the theorem.
Remark. Importantly, the SOC problem (12) generally has an intractable prior distribution making the direct sampling from μ non-trivial. In Section 3.2, we propose a practical method to obtain a sample from μ.
Interpretation of Theorem 3.1 In this paragraph, highlight special case in which becomes tractable. This case reveals an insightful connection to Denoising Diffusion Samplers (DDS) (Vargas et al., 2023). Since Uo is prescribed and often tractable, μ x -Uo(=)-Vo(aa tractable when (x( is constant. This can be achieved by setting the memoryless reference dynamics from 7) where Ule(o.=) := Uy for some 7 < 1, which ensures that the uncontrolled dynamics are sufficiently mixed between E [0, r]. particular, we choose T = our SOC problem (12) reduces to DDS (Vargas et al., 2023):
5

---

Corollary 3.2. When (1E Up(x)
||=|" 25
and Br o2 the problem (12) degenerates
1/- [ + U1(1)) - -Uo(x))
vo(xi)
min E
(16)
31 X₁ /. Beu dt di - Xo 6o-4,
s.t. dX.,
and setting e Uo and Bg to memoryless noise schedule recovers DDS formulation in eq. (9).
Because ())( (r), the annealing mechanism is absent in DDS, implying that the reference dynamics does not guide the particles toward the desired target distribution v.
3.2 Debiasing Annealed Adjoint Sampler through Learning Prior
From now on, we consider the specific case where the initial potential given by Uo(z) ||
2 i.e.
eto xx (N0, T)). We begin by formulating a stochastic optimal control (SOC) problem whose optimal dynamics generates samples Xo from the (intractable) prior distribution μr.
Lemma 3.3. Let Un(x)=
- Given the value function V : X x 10, 1] IR of the SOC problem
2
eq. (12), consider the following SOC problem:
min [ (x)) Vo(Xo)
(17)
s.t. dX, (X,) odW,, X-1 6g-
Note that the time interval of the underlying dynamics is [-1, 0). Then, the terminal distribution induced by optimal control dynamics is μ -oo())-Vo())
By the definition of the value function in eq. (3), Vo (2) is the expected cost of the SOC problem (12)
conditioned by Xo x. Thus, V-(-) in eq. (17) can replaced by eq. (12). As result, arrive a unified SOC problem that defines an unbiased annealed sampling framework:
Theorem 3.4. Consider the following SOC problem:
[-] [x))( (.' (u))(x)) [ ('
min E 0,8
(18)
d.X, (X.)dl + odW. for - X da,
(19)
of dx,
vU,(X)) + (x)) dt o,dW, for 0 0st5l.
(20)
2
Then, the optimal control dynamics is an unbiased sampler, ie., pi x uU
3.3 Practical Implementation
In this section, we present a practical algorithm for solving the SOC problem eq. (18). Our approach alternates between updating the control policies and for the annealed dynamics (20) and the prior-estimation dynamics (19), respectively. We employ adjoint matching to gain a scalability and efficiency. The following two paragraphs describes the training objective for the controls and
Adjoint Matching for Annealed Dynamics We now describe how to optimize the control function (z( for annealed control dynamics eq. (20). Let u" be the parameterization of this control. Given an initial state Xo sampled from eq. (19), we simulate the forward trajectory = {X, ]ee(o.1) using the current control estimate " To update the control u we adapt the AM framework eq. (4), yielding
((u,:)) = /'' (X,) + ora(t; X)l X X~p", p", á = stopgrad(u")
(21)
where da(t; dt l((t:)) (- + avv.(()))
a(l;X) = 0,
(22)

---

Algorithm 1 Non-equilibrium Annealed Adjoint Sampler (NAAS)
Require: Tractable energy Un. twice-differentiatl target energy U1 (z), two parametrized networks u² [] x X + X' and v [-1,0]] x X' + X' for control parametrization, buffers B, and B Number of epochs N₂ and N..
1: for stage k in 1, 2.... do
2: for epoch in 1, 1........ do
3: Sample (Xe [te through eq. (19) and eq. (20).
forward pass > backward pass D add to buffer
4: Compute lean adjoint (ar through eq. (22).
5: Push {(t, 0,, Xe))ce(o.1] into a buffer B,.
6: Optimize u" by eq. (21) with samples from B...
b adjoint matching
7: end for
8:
for epoch in 1. do
9:
Sample (et)<(_1.1]] through eq. (19) and eq. (20).
forward pass > backward pass add to buffer
Compute lean adjoint (erj)) through eq. (22).
Push (Xo, an) into a buffer B,..
12: Optimize e by eq. (24) and (25) with samples from B,
reciprocal adjoint matching
13: end for
14:
end for
where Ue is the Hessian of Uz) whose vector- vector- Hessian product can be computed efficiently by
a(t; x) ( of (PU((X)) = of (a(t:
(-2
ä stopgrad(a).
(23)
Note that this AM scheme is not computationally efficient AS (Havens et al., 2025), since it requires solving a backward ODE due to the use of annealed reference dynamics. Nevertheless, remains more efficient than alternative approaches such as naive backpropagation or the Bor-iariance method proposed in Richter and Berner (2024).
Reciprocal Adjoint Matching for Sampling the Prior Similarly, we address the optimization of the control function (ar) corresponding to eq. (19), parameterized as v., We employ the reciprocal AM framework learn v.. Specifically, we begin by sampling the initial states Xo from controlled dynamics with current control v². We then generate full trajectories X = {X+)<e(o.1] using the annealed dynamics associated with u", and simulate the backward adjoint process via eq. (22). Then, we obtain a state adjoint pair (Xn, 4n) on f O. Since the reference dynamics w.r.t. v has vanishing drift and state cost, apply reciprocal adjoint matching (Havens al., 2025) follows:
c(u"; = [" (x) + a
stopgrad(u),
(24)
where X, ~ Prjo (- :- /' ((1 (1++))Xo (1 + +)2 -))22)))
(25)
Thus, to update v is only required to cache state-adjoint pair (Xo, an) at t = 0, which is memory efficient. For the detailed derivation of conditional distribution, Appendix A.
system at every iteration is computationally expensive. To address this, we employ replay buffers that store the necessary state-adjoint pairs used train the control networks. Specifically, instead of recomputing the lean adjoint system each step, we refresh the buffer every 200-500 iterations, which significantly reduces computational overhead. For training u², we maintain buffer B,, containing triplets (f, Xx, az) for t E [0, 1]. For training vº we use separate buffer B, to store the state-adjoint pair (Xo, ao) at t 0. The implementation details are provided in Appendix D.
Replay Buffers Computing the full
Algorithm As described in Algorithm 1, we optimize the control functions u and v" in an alternating fashion using the adjoint matching objective desribed in eq. (21) and eq. (24). To compute the required state-adjoint pair, we solve the forward-backward systems defined in eq. (22) and eq. (25). To avoid the computational overhead of solving these systems at every iteration, we adopt a replay buffer strategy that caches previously computed state-adjoint pairs.
7

---

Groundtruth
NAAS (ours)
Groundtruth
NAAS (ours)
2213
3000
(a) GMM40 (50d)
(b) MoS (50d)
Figure 2: Qualitative Comparison on GMM40 (50d) and MoS (50d). For each task, we visualize the kernel density estimate (KDE) of the target distribution alongside the generated samples (shown as red dots), projected onto the first two principal axes. The samples from NAAS closely follow the structure and support of the ground-truth, demonstrating accurate mode coverage and high sample fidelity in high-dimensional settings.
The role of two optimization stages The two optimization problems in our framework, (i) learning the prior distribution f over the interval [-1, 0], and (ii) optimizing the controlled annealed dynamics over [0, 1], exhibit fundamentally different learning dynamics and objectives.
The first-stage optimization over [-1, learns prior distribution / e-Vo-Vo where Un is simple, prescribed energy function (typically corresponding to Gaussian distribution). This prior designed to have broad and smooth support, in contrast to the often sharp multimodal structure the target distribution v. The objective here is not to capture fine-scale structure, but to construct a low-complexity distribution that can be easily sampled using simple base process, while still broadly covering the support of
contrast, the second-stage stage optimization over [0, 1] focuses on refining the annealed dynamics to guide trajectories from toward the target distribution v. Since the annealed reference dynamics already bias the flow toward high-density regions of the target, the control this stage primarily serves to correct residual mismatches. Because we only required to learn the residual mismatch, the second optimization problem also easy to learn.
By decoupling these two objectives, the alternating scheme enables more stable and effective training. Each optimization stage solves simpler and more focused subproblem: the first constructs well-behaved initialization, and the second incrementally corrects the flow toward the target. This separation reduces optimization interference and improves both convergence and sample quality.
4 Related Works
Sequential Monte Carlo and Annealed Sampler Many recent works have focused on Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) methods (Chopin, 2002; Del Moral et al., 2006; Guo et al., 2024), which sequentially sample from a series of intermediate (often annealed) distributions that bridge tractable prior and complex target distribution. A large number of these methods leverage annealed importance sampling (AIS) (Neal, 2001; Guo et al., 2025) to correct for discrepancies between the proposal and target distributions, often improving the quality of proposals over time. Several earlier approaches that do not rely on neural networks (Wu et al., 2020; Geffner and Domke, 2021; Heng et al., 2020) have been developed to learn transition kernels for MCMC. Later, transition kernels has been explicitly parameterized using normalizing flows (Dinh et al., 2014; Rezende and Mohamed, 2015), enabling more expressive and tractable models. Based on this idea, various approaches (Arbel et al., 2021; Matthews et al., 2022) have been proposed by refining SMC proposals via variational inference with normalizing flows. More recently, diffusion- -bsed annealed samplers (Vargas et al., 2024; Aahound-Sadegl et al., 2024; Phillips et al., 2024), inspired by ccore-based generative modeling (Song et al., 2021). have gained attention. These methods utilize annealed dynamics derived from denoising diffusion processes and typically combine importance sampling with a matching loss to account for the mismatch between the sampling trajectory and the target distribution. On other line of work, Sequential Control for Langevin Dynamics (SCLD) (Chen et al., 2024) improves training efficiency by employing off ooo-policy optimization over annealed dynamics. However, it still relies on importance sampling to ensure accurate marginal sampling.
8

---

Table 2: Results on synthetic energy functions. We report Sinkhorn distance ()) and MMD (L).
References for each comparison method are provided in Table I and Section 5. All comparison values are taken from SCLD (Chen et al., 2024).
MW54 (d
Funnel 10)
GMM40 (d 50)
MoS (d
Method
Sinkhor
MMD
Sinkhorn
Sinkhorn
MMD
Sinkhorn
SMC
20.71 + 5.33
149.35 4.73
46370.34 + 137.79
3297.28 + 2184.54
SMC-ESS
=
117.48 : 9.70
24240:68 50.52
1477.04 133.80
CRAFT
11.47 0.90
0.115 = 0.003
134.34 0.66
28960.70 + 354.89
0.257 = 0.024
1918.14 108.22
DDS
0.63 +: 0.24
0.172 + 0.031
9.55 142.89
5435.18 172.20
0.131 +: 0.001
2154.88 # 3.861
PIS
0.42 + 0.01
10405.75 + 69.41
0.218 + 0.007
2113.17 # 31.17
AS
0.32 = 0.06
18984,21 62.12
0.210 = 0.004
2178.60 = 54.82
CMCD-KI.
0.57 0.05
0.095 + 0.003
513.33 + 192.4
22132.28 + 595.18
1848.89 532.56
CMCD-LV
0.51 = 0.08
139.07 9.35
4258.57
737.15
1945.71 = 48.79
SCLD
+
134.23
8.39
3787.73
249.75
656. 10 88.97
NAAS (Ours)
0.10 +00.0075 0.076 +: 0.004
132.30 + 5.87
496.48 :: 27.08
0.113 + 0.070
394.55 +: 29.35
Experiments
In this section, we evaluate NAAS across diverse set of sampling benchmarks. Further implementa- tion details are provided in Appendix D.
Our primary focus challenging synthetic distributions, including the 32-mode many-well (MW54, 5d) distribution, the 10d Funnel benchmark, a 40-mode Gaussian mixture model (GMM40) in 50d, and 50d Student mixture model (MoS). We assess the performance using the maximum mean discrepancy (MMD) following Akhound- Schound-Sadegh et al. (2024) and the Sinkhorn distance (Sinkhomn) with small entropic regularization (10 33) following Chen et al. (2024).
We demonstrate the practical upplicability of NAAS by extending our evaluation to the generation of Alanine Dipeptide (AD) molecular structures. This molecule consists of 22 atoms in 3D. Following the implementation of Midgley et al. (2023); Wu et al. (2020), we use the OpenMM library (Fastman et al., 2017) and represent the molecular structure using internal coordinates, resulting
in 60-dimensional state space. We evaluate KI. divergence (DxL) between 2000 generated and reference samples for the backbone angles Φ, & and the methyl rotation angles 71. 72, 73. which are referred to as torsions. Moreover, report Wasserstein distance between energy values of generated and reference samples (2).
As shown in Table 2, our proposed method, NAAS, consistently outperforms existing baselines across the majority of synthetic sampling benchmarks. On the MW554 and GMM40 tasks, NAAS improves over the second-bes best method by more than 75%, demonstrating exceptionally strong performance. Furthermore, on both the Funnel (MMD) and MoS benchmarks. NAAS achieves sub- stantial improvements over competing methods. These results collectively highlight the effectiveness of NAAS as a high-quality sampler across diverse and challenging targets.
Main Results
Figure 2 presents qualitative results for the GMM40 and MoS benchmarks, showing side- dide-by-side visualizations of samples generated by NAAS and the corresponding ground truth. The generated samples accurately capture the structure and support of the true distributions. All modes are well- represented without collapse or over-concentration, indicating high sample diversity and fidelity.
Figure illustrates the training dynamics of NAAS, with both MMD and Sinkhorn distance plotted over training iterations. In particular, the MoS experiment Figures 3b and 3c highlight that our initial dynamics are remarkably strong; our model begins training with an MMD around 0.22 and a Sinkhorn distance near 2600. These values are already competitive with, or even superior to, the final performance achieved by several established baselines such as SMC, SMC-ESS (Buchholz et al., 2021), PIS (Zhang and Chen, 2022), and DDS (Vargas et al., 2023). All comparison benchmarks. including SMC, SMC-ESS, and CMCD (Vargas et al., 2024) are taken from SCLD (Chen et al., 2024). This favorable starting point is direct consequence of the annealed reference dynamics employed by our method, which provides a meaningful initialization that captures coarse features of the target distribution. NAAS then further improves the sampler by learning the control. Ultimately, NAAS surpasses the second-ber best method, SCLD (Chen et al., 2024), by more than 30% in both metrics,

---

mamber 2000s
Narter of Epedto
Number of tpec
(a) Funnel (MMD)
(b) MoS (MMD)
(c) MoS (Sinkhorn)
Figure 3: Visualization of training dynamics over epochs. Each are presented in different color.
demonstrating that its performance gains are not merely due to favorable initialization, but also to its ability to effectively learn corrections the reference dynamics. This highlights the strength of NAAS in combining informative priors with principled learning for efficient and high-fidelity sampling.
Alanine Dipeptide Generation We compare our model to PIS (Zhang and Chen, 2022) and DDS (Vargas al., 2023), which can be viewed special case of framework that excludes the annealed path measure. As shown in Table 3, our model achieves significantly better performance in terms of ?1 and E() while showing comparable results the remaining metries. It also achieves favorable values across all torsion angles while maintaining similar energy distribution, indicating the absence of mode collapse. Furthermore, illustrated by the qualitative results in Figure 4, the torsion angle distributions produced by our model closely follow the trends of the reference density.
Table 3: Results for AD generation. We report
KI. divergence for torsions (DR.)) and Wasserstein 2-distance for energies F22)
Method
71
72
18
E IW2
PIS
0.597
0.952
0.453
0.498
7,038
5,918
Comparison of five torsions between generated and reference samples
Figure 4:
0.150
NAAS (Ours)
0.250
0.236
0.272
0.132
0.156
1.076
6 Conclusion
We proposed NAAS, novel and unbiased sampling algorithm grounded in SOC theory under annealed reference dynamics. Our key contribution is on formulating the SOC problem that results unbiased diffusion sampler tailored to have annealed reference process. To construct an unbiased estimator, we propose joint SOC problem involving two distinct control functions, which are optimized in an alternating manner using an adjoint matching. We validated NAAS on diverse set of synthetic benchmarks, demonstrating state-of-the-ar performance across all tasks. Additionally, we extended our framework molecular structure generation with alanine dipeptide, suggesting that NAAS can scale to more complex and structured domains. Despite these encouraging results, it remains be tested its effectiveness on larger-scale and real- world scenarios, potentially with more expressive annealing schedules. Moreover, deeper investigation into the trade-off between the two optimization enaess-xammning how the strength the annealing term influences the behavior of NAAS-remaing important direction for future work. As limitation, NAAS requires computing the Hessian the given energy function during sample generation, which introduces computational overhead. Moreover, NAAS requires two networks to train due to its alternating scheme. We do not see any negative socictal consequences of our work that should be highlighted here.
Acknowledgments and Disclosure of Funding
JC and YC acknowledge support from NSF Grants 2450378. MT thankful for partial supports by NSF Grants DMS-1847802, DMS-2513699 DOE Grants NA0004261, SC0026274, and Richard Duke Fellowship.
DMS-2206576, and CMMI-
10

---

References
Tara Akhound-Sadegh. Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marein Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, al. Iterated denoising energy matching for sampling from boltzmann densities. In Proceedings of the 41st International Conference on Machine Learning, pages 760-786, 2024.
Michael S Albergo and Eric Vanden- Rijnen. Nets: A non-equilibrium transport sampler. arXiv preprint arXiv:2410.02711, 2024.
Michael Arbel, Alex Matthews, and Arnaud Doucet. Annealed flow transport monte carlo, In International Conference on Machine Learning, pages 318-330, PMLR, 2021.
Hamidreza Behjoo and Michael Chertkov. Harmonic path integral diffusion. IEEE Access, 2025.
Richard Bellman. The theory of dynamic programming. Technical report, Rand corp santa monica ca, 1954.
Kurt Binder, Dieter W Heermann, and K Binder. Monte Carlo simulation in statistical physics, volume 8. Springer, 1992.
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond elbos: a large-scale evaluation of variational methods for sampling. In Proceedings of the 41st International Conference Machine Learning, pages
2024.
Joey Bose, Tara Akhound- Sadegh, Guillaume Huguet, Kilian FATRAS, Jarrid Rector- Bettor-rrook,, Cheng- Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael M Bronstein, and Alexander Tong. Se (3)-stochastic flow matching for protein backbone generation. In The Twelfth International Conference Learning Representations, 2024.
George EP Box and George C Tiao. Bayesian inference in statistical analysis. John Wiley & Sons, 2011.
Alexander Buchholz, Nicolas Chopin, and Pierre E Jacob. Adaptive tuning of hamiltonian monte carlo within sequential monte carlo. Bayesian Analysis, 16(3):745-771, 2021.
Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anand- kumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.07081, 2024.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrödinger bridge. SIAM Review, 63(2):249-311 2021.
Nicolas Chopin. A sequential particle filter method for static models. Biometrika, 8933)533-5552 2002.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-329, 1991.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411 1-436, 2006.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components
estimation. arXiv preprint arXiv:1410.851P 16, 2014.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference on Learning Representations (ICLR), 2025.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development high performance algorithms for molecular dynamics. PLoS computational biology, 13(7):e 1005659, 2017.

---

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel. Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021. attp://jmlzz org/ wwwrgg/paps////aee pwpppp//22/20--e/eeeppp 2221222222222222 html.
Tomas Geffner and Justin Domke. Mcme variational inference via uncorrected hamiltonian annealing, Advances in Neural Information Processing Systems, 34:639-651, 2021.
Wei Guo, Molei Tao, and Yongxin Chen. Provable benefit of annealed langevin monte carlo for non-log-concave sampling. arXiv preprint arXiv:2407. 16936, 2024.
Wei Guo, Molei Tao, and Yongxin Chen. Complexity analysis of normalizing constant estima- tion: from jarzynski equality to annealed importance sampling and beyond. arXiv preprint arXiv:2502.04575, 2025.
Aaron Havens, Benjamin Kurt Miller, Bing Yan, Carles Domingo-Enrieh Anuroop Sriram, Brandon Wood, Daniel Levine, Bin Hu, Brandon Amos, Brian Karrer. et al. Adjoint sampling: Highly scalable diffusion samplers via adjoint matching. arXiv preprint arXiv22544117733, 2025.
Jeremy Heng, Adrian N Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequential monte carlo. The Annals of Statistics, 48(5):2904-292)) 2020.
Eberhard Hopf. The partial differential equation ut+ uux= μxx. Communications on Pure and Applied mathematics, 3(3):201-233) 1950.
Jarzyssii Nonequilibrium equality for free energy differences. Physical Review Letters, 78(14):2690, 1997.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-558) 2021.
Hilbert Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of Statistical Mechanics: Theory and Experiment, 2005 1):11 2005.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Jean-François Le Gall. Brownian motion, martingales, and stochastic calculus. Springer, 2016.
Christian Léonard, Sylvie Rælly, and Jean-Claude Zambrini. Reciprocal processes. measure- theoretical point of view. Probability Surveys, 2014.
Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos A Theodorou. Deep generalized Schrödinger bridge. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky TQ Chen. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations, 2024.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Amaud Doucet. Continual repeated annealed flow transport monte carlo. In International Conference Machine Learning, pages 15196-15219. PMLR, 2022.
Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation of state calculations by fast computing machines. The journal of chemical physics, 1186))108711092, 1953.
12

---

Laurence Illing Midgley, Vincent Stimper. Gregor NC Simm, Bernhard Schölkopf. and José Miguel Mernández-Lobard Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023.
Radford M Neal. Annealed importance sampling, Statistics and computing, 1 11:125-139, 2001.
Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767, 2003.
Maxence Noble, Louis Grenioux, Marylou Gabrié, and Alain Oliviero Durmus. Learned reference- based diffusion sampling for multi-modal distributions. In International Conference on Learning Representations (ICLR), 2025.
Nikolas Nüsken and Lorenz Richter. Solving high-dimensional mamiltonnjacobbe jacobi-bellaa bellmann pdes using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial differential equations and applications, 2(4):48, 2021.
Michele Pavon and Anton Wakolbinger. On free energy, stochastic control, and schrödinger processes. In Modeling. Estimation and Control of Systems with Uncertainty, pages 334-348 Springer, 1991.
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. Particle denoising diffusion sampler. arXiv preprint arXiv:2402.06322 2024.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMLR, 2015.
Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In The Twelfth International Conference Learning Representations, 2024.
Hannes Risken. Fokker-Planck equation. In The Fokker-Planck Equation, pages 63-95. Springer, 1996.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
Gabriel Stoltz, Mathias Rousset, et Free energy computations: mathematical perspective. World Scientific, 2010.
Emanuel Todorov. Linearly-solvable Markov decision problems. In Advances in Neural Information Processing Systems (NeurIPS), 2007.
Mark E Tuckerman. Statistical mechanics: theory and molecular simulation. Oxford university press, 2023.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In International Conference on Learning Representations (ICLR), 2023.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets varia- tional inference: Controlled monte carlo diffusions. In International Conference on Learning Representations (ICLR), 2024.
David Weininger. Smiles, chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36, 1988.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in neural information processing systems, 33:5933-5944, 2020.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. In International Conference on Learning Representations (ICLR), 2022.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Our abstract and introduction clearly state the claim including contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification:
We provide assumptions and proofs in Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all details necessary reproduce the results in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [No]
Justification: All data used in this work is open-source. Due to organizational policy, we are currently unable release source code, but plan to make it publicly available as soon as possible.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (uttps://nips.. cc/
Saaides/CodeSubiisee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results, See the NeurIPS code and data submission guidelines (https:
//nips. cttpppsccppbbcccccccc httliic/guiiiee guwdes/ Codeeubmie wwwibbmmsssiinoo asiionolicyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: See Table 2. We report mean and variance of each experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
16

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: httpss///aarpps...ppbbb///////// wc/ppblic/Ethttccccccccc cwchchhescceeellecc ines?
Answer: [Yes]
Justification:
We read and comply with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Section 6.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the original papers of the authors that release the assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not utilize crowdsourcing or do research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not utilize crowdsourcing or do research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this paper does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https battps://ee /ttp://eeeuu.pp/eeee 2ttp://p////..oeeeeeee////////// for what should or should not be described.
20

---

A Proofs and Additional Theorems
A.1 Proof of Theorem 3.1
Proof. Let V : [0, 1] x & - R be a value function defined as eq. (3). Then, the necessary condition,
namely Mamniton-Jacob--elllnee
(HJB) equation (Bellman, 1954), is written follows:
of
of 2 AVi(2) of [VV + (,U(()), (a = 0,
0 (.v VU.(z) VV((2)
(26)
where (x) = oo (x)... Moreover, the corresponding Tokker-Planchs equation (FPE) (Risken, 1996) can be written as follows:
Viia)) of Oepi(a) v - o(VV)(t) p2(x)) + v2. ( (pille))
(27)
Now, let (x)()) = e Vi(z) and (x(x) = pt (())) Then,
@epr(x)) = 0. (zr)ee Va(a))
o. VUl(T) VV((2) + ot AVi()) - oB 2 ( e eVicoe, (28)
2
2
VU(()) Vpr(z) - of
(z) + (,U((t))))()..
2
2
Following the eq. (52) in Liu et al. (2022), we can derive
Opt
0V,
0g0: exp(V))
Pt
01
at.
ot
AoAVVP + VVV. Vp + of App
cxp(Vi)
(v
VUPt
2
of 2
AVi + of of 2
+ exx((V))p
VU,
VV2
V 2 - 0,U,
2
o²
exp(Vi)
(2AViPt + 2VV. Vpt - Apt + VV AV)
2
+
exp(VV) (V (VU - -0,0, exp (Vi)tt
2
v.(VU19.)
a² 2
+ 0²
AQi
V (VU())) - 0,U0..
2
Therefore,
(x) of v (VU,4 V&..) 0.U,
Therefore, by combining these equation, we obtain joint PDE, namely Hopf-Cole transform (Hopf, 1950), i.e.
a² VU,- V - + 2 (x))(x) of V (VU + V - 0,0,i (z)))x(()) v(z). (x)
Ap. 0,00P po(z)(o(z) plx),
(29)
Notice that (x) (z) and yp (z) constant by the definition. Morcover, by simple calculation, one can derive that & has the following closed-form expression:
(x(x)) Ce (U.(2)
up some multiplicative constant C. Combining these equations obtain (x(x) (x) xe (1(2) (z()) = exp(-U,(z)- -v))))
Hence,
pz(x) x (xp(-UU(x)-Vo((a)) (x) x
CC (ex)(-U((z)) x v(x). (30)
21

---

A.2 Proof of Corollary 3.2
Proof. Let (x)) Uo(2)
20 and 3r = Then,
o v(( o (U)
Be z.
2
=2"
Thus, by substituting all these variables into eq. (12), the equation degenerates to
[/" (llue(X)))l + U1(1)) -
min E
(31)
31
s.t. d.X.
Xt + /. (x)) dt + / Xo ~ u.
2
By Feynman Kac formula (Le Gall, 2016), exp (-Vo(Xo)
E Ex ( (x)) /exp (-(Ui(X,) Uo(Xi)))
(32)
eq. E Xx (exp (exp(-(04) (-(Ui(X;) U6(Xi))) constant.
(33)
Therefore, by substituting Vo constant to eq. (30), obtain
(z(x) C exp(-U U,) NN
(34)
for some constant C.
Note that since
base (X1)Xo) =N (xi:ee= Bid Xo. 0² 20 6 )))
(35)
the reference dynamics is (almost) memoryless when So 3, co.
A.3 Proof of Lemma 3.3
Proof. Since (u(x)) x e Un(a)
Vo(X.) log Åo (Xo) log (Xx) log 1og,N(1;0 - -log
constant,
the SOC problem (17) can be rewritten as follows:
[[] (ll Xx,
0
log N(Xoi0,
min E 16
(x, 1)1
(((g)
(36)
dX, ====u(x,, (P,(xx,)ddt + adW, X_1 6g-
By setting the terminal target distribution as μu, this optimization problem becomes identical to the formulation in eq. (8), which refer to as PIS. Consequently, the terminal distribution induced by its optimal solution is guaranteed to match μ.
A.4 Proof of Theorem 3.4
Proof. Let V : [-1, x X + IR be value function defined on the given SOC problem (18). Then, the corresponding HJB equation can be written as follows:
σ²
0 = a. (.V((a))+++
2 AVi(x)
[V)(2)) for < 0, eq. (26) for t 2 0.
(37)
Let (ar) = e (v(w and (p(x) = p (z() for -1 < S 1. Then, similar to the proof of Theorem 3.1, one can easily derive the following paired PDE:
01 = -e² Asp
-
0,g(x,t) 2 A
a²
when t t< 0,
plr, 1) = 1,
(38)
(0p((,t) 3. VVU Vip d +A++0,Up
when t 2 0,
(39)
0,p(2,t) -V.
PVV(UU + v() - 0,U,0
22

---

As shown in the proof of Theorem 3.1, one can check that the solution (for some multiplicative constant C) eq. (39)
(x(x) = Ce U.0,, t 120 0.
(40)
Therefore, Po(x) Ca-Ua(a) for some multiplicative constant C.
Now, over the interval E 1- 0], the function & evolves according to the heat equation. Thus, we can explicitly write the solution for on [-1, 0] given the terminal condition eq. (40) follows:
I=l² (x(x) C'e 204 1J for - < <<
for some multiplicative constant C². Note that lim,-- Sr. ög. Combining these equation, ; admits explicit solution, which can be written in a piecewise form over the full domain:
Hell
(x(x) C'è
20 Ce-U(())
for < <,
(41)
for 0 0115l. 1.
Thanks to the explicitly written Ge., we can easily derive the following conditions:
f-1Ộ-1 = 60. Pufo x e -V-o% x μ4,
00 & e- x v.
(42)
Therefore, the solution to our stochastic optimal control (SOC) formulation in (18) yields the desired unbiased sampler.
A.5 Detailed Explanation of Reciprocal Adjoint Matching
By using the lean adjoint matching loss, update both u", v.. Precisely, the original lean adjoint matching loss can be written as follows:
C(u v x) [9] |/ (x,) oa(t: dt + L.' [x,)(X)) + t, dt,
(43)
where
X ~ ((,00,
(),ū) stopgrad((r* (o", =0
(44)
da(1; X) dt
x(t;:) (- a.VU,(x))
t t2 0,
(45)
da X) 0, 1 <<< 0.
(46)
di Thus, in the backward interval t < 0, the lean adjoint variable remains constant: a(t; = a(0; X) =: ag for all t < O. Furthermore, as discussed in eq. (6), the optimal backward control 1 satisfies the reciprocal property. That is, under the optimal dynamics, the marginal distribution of X, given the endpoints X_1 = and Xo corresponds to a conditional Gaussian:
X₁
P X-1 == Xo) = X, ~ N ((++ 1)xo, (t 1) - 1))
base
(47)
for all e [-1, 0). By leveraging this optimality condition, we obtain our reciprocal adjoint matching objective as eq. (24).
B Further Discussion on NAAS
B.1 Discussion on Scalability and Algorithmic Stability
This section elaborates on three aspects NAAS: (i) scalability with two learnable control functions, (ii) the rationale behind the alternating optimization scheme, and (iii) the complexity and stability of Algorithm I.
Scalability Although NAAS employs two control networks, the overall computational and memory costs remain comparable existing neural control frameworks, Both control functions share the same lightweight architecture and are trained over disjoint annealing intervals, such that the doubled parameter count does not lead to doubled runtime or memory overhead. Empirically, both stages converge within similar iteration counts, and the per-iteration training time remains nearly unchanged since gradients are backpropagated through independent adjoint integrations. This modular design preserves scalability even in high-dimensional settings.
23

---

Alternating Optimization The optimization of NAAS is deliberately decoupled into two stages, each targeting a distinct subproblem:
Stage (t E [-1, Ol) learns smooth, broad prior distribution that provides wide support over the target space. This stage functions as a coarse-scale initialization similar multiscale optimization methods.
Stage (i E [0, 1) refines the control to transport samples from the coarse prior to the target distribution. As the annealed dynamics naturally bias trajectories toward high-density regions, this refinement primarily corrects residual mismatches, ensuring stability and data efficiency.
This alternating design mitigates interference between global and local objectives, leading to more stable convergence than a fully joint optimization. It also facilitates clearer diagnostics, as each stage's broe-broad coverage versus fine refinement-can be independently verified.
Empirical Comparison with Joint Optimization To verify the effect of our alternating algorithm, we compared it with joint optimization scheme that trains single control over the entire time horizon E 1,. Evaluations were conducted on the MW54 and MoS benchmarks. As shown in Table 4, our alternating optimization consistently outperforms the joint optimization across both tasks. These empirical findings align well with our theoretical motivation, demonstrating that decomposing the learning process into two focused stages yields better convergence and overall performance.
Table 4: Comparison between joint and alternating optimization schemes on MW54 and MoS. We report Sinkhorn distance (L).
Scheme
MW54
MoS
Joint
0.20
597.99
Alternating (Ours)
0.10
394.55
B.2 Difference between NAAS and NETS
Non-equilibrium Sampling In a similar spirit to AIS (Neal, 2001), non-equilibrium samplers are methods that progressively guide samples over finite horizon using meaningful reference annealing path defined by the annealing energy U, between target and prior distributions, for instance, the same reference path eq. (1 l) is considered for both NAAS and NETS (Albergo and Wanden-Bijndee 2024). Such guidance from non-equilibrium or reference distribution exp(-Up indeed plays key role in these methods.
Difference between NETS and NAAS NETS (Albergo and Vanden Mijndnn, 2024) apply Jarzyn- ski's equality (Jarzynski, 1997) to match intermediate distribution at time to be proportional to exp(-U+). So, in NETS, an additional component u, in drift term is introduced to eliminate the discrepancy between reference dynamics and the desired distribution pt x exp(-U.).
On the other hand, while NAAS also appends an additional term u, to the same reference path eq. (11),, its de is derived within the stochastic optimal control framework, aiming to solve specific objective formulated in eq. (12). This distinction alters the effect of u, in modifying the reference path eq. (11) between the two methods, the objective of NAAS eq. (12) does not explicitly enforce u to match the reference distribution exp( UU)) Indeed, we prove in Section 3.1 that the proposed by NAAS results in a controlled SDE whose instantaneous distribution Pt satisfies Pt x exp(-U, VT). where Vi 0. That is, the instantaneous distribution of NAAS is highly correlated with the reference distribution exp(-U,). but biased by an additional term exp( exp(-Vt),, which vanishes at the terminal time / = 1. Furthermore, we highlight that since Vi 0. Pi ε exp(-Ut
Vi) exp(-Ui) x Hence, the controlled SDES proposed by NAAS do converge to the desired target distribution at terminal time t = 1.
24

---

B.3 Connection between NAAS and PDDS
Particle Denoising Diffusion Sampler (PDDS) (Phillips et al., 2024) considers stochastic process
between Normal distribution / and the target 2
dy, = 1-
dBB. No,N( ~~-== T 3. (48)
where the "guidance" ge given by
(v) = / (0(2))o((/22)= = y)dx,
90 N
(49)
-2 ( 3,d7)
p(agaae) x1) = NN v Ax1,
= exp
(50)
and can be learned, c(z, 1) RJ V log 91- (z), via the NSM loss:
! El ux, 1;0 - -1 v log go (xo)
(-fo B-dT)
R4 Ky== 1 -===i->==a exp
(51)
where the expectation is taken over the optimal distribution corrected via IWS.
SOC formulation for PDDS In this paragraph, we introduce the SOC problem that corresponds to the optimal solution of PDDS. Moreover, further discuss that the adjoint matching loss for the SOC problem is equivalent to NSM loss.
Proposition B.1. Equation (48) is the optimal controlled process to the following SOC problem
1 2 (x,)(())))22 + log 3 (xx)
min E
(52)
s.t. dX,- = -- + v 2/ (2 dt + ddB. Xo ~o~N(0,
(53)
That is, the optimal control (52)
(x)x) = / V (ogg)-((x))
(54)
Furthermore, the lean adjoint matching loss (W.E.. the optimal distribution) for (52) is the same as the NSM loss (51).
Proof. The optimal control to (52) follows
ui (x) =
- (z)
--V23,-V2 (- lllog E Ex.gxx-moa -m ( exxp (-log log N (x₂)
2 (10g] / N (01 = 2)d21
2 log
where the last equality follows by the fact that pr(2(x.)) = -N(x;; /1 X A A1e- due to Pt = I) for all t e [0, 1]. and hence
ʃ AwDiilll r)dy - ff 9o00))1-22(00011- xddy g1-r(x).
Next, notice that the lean adjoint ODE in this case reads
dā;
-rs ō
V log . (Xz).
(55)
dt
We substitute the identity, V log (z) + x V log g. (x), into the original NSM loss, E|] V log n. (X.)
+
X. Ks V log (
25

---

The solution to (55) admits the form ār C exp
where C' such that ā
Cexp (So 31- dl-) V log 7 Hence, (55) can solved closed form
ā Vlog N. exp (fo B1 31-,ds
['' B1-
7
[ 31-ad8
V Voog" N'
7 exp
V log
N 7
exp
3,
loog 9o R1-ts
where the third equality is due to change of variable := 1 Applying lean adjoint matching,
Ellu(a, 0) + / 22 = Bl|u([, 1; /2 Ki V log go
(56)
with specific parametrization u(x, t: 2 v(z 0) yields the NSM loss (51).
Connection to NAAS The stochastic optimal control (SOC) problem underlying PDDS (Phillips et al.. 2024)-formulated in eq. (52)-shares the same structure as the SOC problem defined in DDS (Vargas et al., 2023), i.e. eq. (9). As discussed in Section 3.1, the DDS objective itself special case of the more general SOC formulation introduced by NAAS (12). Therefore, the SOC perspective of PDDS naturally inherits this structure and can also be regarded within the broader NAAS framework. Despite the close connection in problem formulation, the methods used to solve the SOC problem are different. NAAS and DDS solve the SOC problem using baased learning objectives. In contrast, PDDS employs a non-adjoint-buss based objective, NSM loss (51). Instead, PDDS use the IWS strategy.
special case
B.4 Solving NAAS via Importance Sampling
In this section, we explore alternative algorithmic approach for optimizing the SOC problem defined eq. (18). While we do not empirically evaluate this method in the current work, we present it promising direction for future research. Our SOC formulation eq. (18) can be addressed in various method, resulting flexible algorithm design. Here, we provide another algorithm that handles SOC problem based on important weighted sampling (IWS). Then, we introduce the close connection to PDDS (Phillips et al., 2024).
Basic formula Consider the SOC problem in eq. (1). In the seminar papers (Dai Pra, 1991; Pavon and Wakolbinger, 1991), the Radon-Nikodym (RN) derivative between the optimal and reference path measures can be written as follows:
dp" da
(X) = exp (-/, (((X))dt (-x()axx)dd-g 9(x,))x(xx))))) Vo(Xo)
(57)
Now, let pd be the path measure induced by the control u. By the Girsanov theorem (see Nüsken and Richter (2021)), we derive
dp'
dppase
(X) dp*
dpu (X)
dpbase (X)
dp
a exp (1,'
- llee xl) - (X.)) dw.) exp (-f- fi(x,)dt -(xi)) /Po(Xoo.
1
(58)
where
(o)))) = e Vo(oa) (x) exp (- (-f- fi(x,)ae si(x:)al- -0xxi]
by Feynman-Kee Kanmman-Kaaaaaa Kac formula (Le Gall, 2016).
26

---

Importance Sampling for NAAS Now, consider our SOC problem (18), Because X-1 0, eq. (58) can be rewritten as follows:
dp" (x) x exp ( 1121lxxx))) (xx
dp'
-x(()) .dW,
eep (1,"
- 1 (l -u (Xe) dW, dW.-J, - S. (xU( (59)
Note that discard ( since X_1 = O. Given this expression, the optimal path distribution p by sampling trajectories from p" and reweighting them using the Radon-Nikodym derivative (59). Specifically, the following importance sampling procedure allows us to recover samples from p":
approximate
For given N, obtain a collection of trajectories { where X ~ p².
2. Compute the importance weights (w where w((1) = # (x by eq. (59).
3. Resample {X Timl
by following categorical distribution:
8 ~~ Cat (1w Ji=1 - (x())
(60)
where (
>w w(t)
Importance Sampled Bridge Matching for NAAS this paragraph, we propose an IWS-based method for learning the optimal control v* over the interval f E []. Since v* governs transition from the Dirac distribution 6n to the intermediate distribution μt, this subproblem can be framed as a Schrödinger bridge problem (SBP) between du and μ.. Specifically, v* is the solution the following stochastic control problem:
min [ (x.)
s.t. dx, (X)dd + odW, X-1 Xo ~ μ. (61)
Recent works (Shi et al., 2023; Liu et al., 2024) introduce the bridge matching (BM) loss as scalable surrogate objective for solving SBPs. Under this framework, optimal control is obtained by minimizing following regression problem:
arg inf Ex [X,, oV (ogpp//Lo,X,,X_X_____ xx)]1
(62)
The only missing component is the ability to sample from the intractable intermediate distribution Given current estimate of the control v², we can (asymptotically) sample from μ using importance- weighted sampling (IWS) as described in eq. (60). By combining IWS with the BM loss, we arrive at an iterative IWS-based algorithm for learning v*, as detailed in Algorithm 2.
C Additional Results
Effect of Learning the Prior Distribution In this section, we demonstrate the effect of learning the prior distribution y in eq. (24). isolate the impact of the reciprocal update, we compare against baseline where o = 0 is held fixed. In other words, we remove the update of o described in lines 8-14 of Algorithm 1. Theoretically, this corresponds to optimizing the SOC problem in eq. (12) under the fixed prior μ -U, which we refer to as a biased estimator. We conduct experiments on a simple wwo-dimensional Gaussian mixture model (GMM) target distribution, denoted GMM-grid. grid, which the target distribution is defined as follows:
3 MΣ
v 9
(1 2.j -0 2),
(63)
Am j=1
As shown Figure 5, both Sinkhorn and MMD measurement of the biased estimator quickly plateaus over training epochs, while the full NAAS algorithm continues to improve, achieving sixfold improvement in performance. Moreover, as illustrated in Figure 6, NAAS provides significantly more uniform mode coverage compared to the biased variant, consistent with theoretical result.
27

---

Algorithm 2 IWS-NAAS
Require: Tractable energy Un. wwice-differentiable target energy U (z), two parametrized networks u² [] x X + X' and v : [-1,0] x X' + X' for control parametrization, buffers B, and B Number of epochs N₂ and N..
1: for stage k in 1, 2.... do
2: for epoch in 1, 1. do
3: Sample (Xe |t(-1,111 through eq. (19) and eq. (20).
forward pass D backward pass D. add to buffer
4: Compute lean adjoint (ar)eto..) through eq. (22).
5: Push {(t, 0,, xe))cefo.1] into a buffer B,-
6: Optimize u" by eq. (21) with samples from B..
b adjoint matching
7:
end for
8:
for epoch in 1. N. do
9:
Sample trajectories {x(1) J,y through eq. (19) and eq. (20).
> forward pass IWS P add to buffer Di bridge matching
Obtain resampled samples ( (x1) Ji=1 through eq. (60).
10:
11:
Push (Xx0) LN Jx=1 into a buffer B..
12: Optimize v by eq. (62) with samples from By.
13: end for
14:
end for
Effect of the Annealed Reference Dynamies We investigate the role of annealed reference dynamics by varying the strength of the annealed reference dynamics. In the annealed SDE defined in eq. (II), stronger guidance can be achieved by increasing the diffusion scale a,t E [0, 1]. Following the geometric noise schedule (Song et al., 2021; Karras et al., 2022), we define ag as
= J o
1-t max V2 12log
max
(64)
min
where min and max are hyperparameters selected by design. Increasing max effectively imposes a stronger guiding signal throughout the training. To assess its impact, we test three values of mmx- 1000, 100, and 10-on the MoS (50d) distribution. As shown in Figure 7, we refer to these settings as strong, moderate, and weak guidance, respectively. As shown in the figure, stronger guidance improves training dynamics: larger max results in lower initial Sinkhorn distance, faster convergence, and superior final performance. These results highlight the critical role of the annealed reference in steering the sampler during optimization.
Results on High-dimensional Sampling Problems We evaluate the scalability of NAAS on a high-dimensional 100D GMM40 task, comparing with two strong baselines: AS (Havens et al., 2025), which employs standard reference dynamics without annealing, and SCLD (Chen et al., 2024), state-of-the-art importance sampling method. We also report wall-clock time for the 50D setting, measured on single NVIDIA A6000 GPU. As shown in Table 5, NAAS achieves superior sample quality and convergence stability, while maintaining competitive runtime. Its efficiency partly stems from the use of replay buffers, which stabilize optimization and enable sample reuse. Although solving the backward adjoint state eq. (22) introduces additional computational overhead, the resulting improvements in stability and sample quality justify this minor cost, highlighting the practicality and scalability of our approach.
Table 5: Results on the GMM40 benchmark across various dimensions.
Method
-
4
d
16
d= 50
d = 100
Running time
SCLD (Chen et al., 2024)
157.90
1033.19
3787.73
12357.49
30-40 mins
AS (Havens et al., 2025)
18984.21
30 mins
NAAS (Ours)
25.15
37.96
496.48
633.20
120 mins
Additional Experiments on Variance Measurement Formally, IWS methods approximate the target distribution by s >N 1i=l (()) 8xta, where the unnormalized weights obtained by w(() oc
28

---

Blas (Sinkhorn)
Unoias (Sinkhom)
Blas (MMD)
Unbias (MMD)
3.18
3.16
Sinkhom
0.8
3.10
0.8
3.08
3.06
-3.0
Epochs
Figure 5: Training dynamies for the GMM-grid energy function over training epochs. We compare NAAS trained without prior correction (Biased, blue) and original NAAS (Unbias, orange). In the original NAAS setup, we update the control e and u in an alternating fashion every epochs. The shaded regions in lightblue and lightcoral indicate periods during which da and u are updated, respectively. Our results show that original NAAS consistently achieves lower Sinkhorn and MMD scores (both 4), indicating that the prior correction term plays key role in debiasing the sampler and improving sample quality.
Groundtruth
Biased
NAAS (ours)
0.0225
0.0000
0.0755
0.0000
Figure 6: Side biddee comparison of ground truth (left), NAAS w/o prior correction (Biased, middle), and NAAS (right). We visualize both the generated samples (red dots) and corresponding kernel density estimates (KDE). As shown, the KDE plot for NAAS (Ours) exhibits more uniform coverage across all modes compared the NAAS w/o correction, demonstrating that learning the prior distribution correctly debiases the sampler.
Strong
Moderate
Weak
Sinkk Sinktom om
3000
2000
1000
200
400
800
800
1000
1200
1800 Number Epochs
1400
1600
2000
2203
2400
2600
2800
3000
Figure 7: Visualization of Sinkhorn metric training epochs for varying levels of reference guidance. We compare three variants of the annealed reference dynamics, each using different guidance strength. Specifically,
set max 1000, 100, and 10 to represent strong, moderate, and weak guidance, respectively. As shown, stronger guidance leads to better initial performance, faster convergence, and superior final performance.
29

---

de (X(()) with the ratio between the target path measure p" and proposal p. When the support or mass of the proposal and target distributions are misaligned- as is common in high-dimensional or complex energy landscapes- these weights can suffer from high variance, as noted in prior works (Del Moral et al., 2006; Stoltz et al., 2010). This leads to poor effective sample size and significant variance in the empirical approximation, and can be further exacerbated when the number of samples N is insufficient.
To see it in empirically, we report the variance of the importance weights across various benchmarks. We choose LV-PIS as the baseline for comparison as the method directly optimizes the variance of importance weights and, similar to our NAAS, starts the generation from Dirac delta. For NAAS, we compute its importance weights using eq. (59). As shown in Table 6, the variance of the importance weights for NAAS is significantly smaller than that of LV-PIS. Remarkably, NAAS achieves this reduction in variance despite never explicitly optimizing the variance of importance weights.
Table 6: Comparison on the Variance of Importance Weights (4)
Method
MW54
Funnel
LV-PIS (Richter and Berner, 2024)
21.11
5,46
NAAS (Ours)
1.84
0.11
Effect of Annealed Reference Dynamics We conduct additional experiments to verify whether the improved performance of NAAS over AS arises from the use of annealing reference. Specifically, we instantiate NAAS with particular choice of U,. which remains equal to Uo for all E [0. 1) and jump U, U₁ at terminal & 1. This setup effectively leads to a two-staged AS baseline that retains the same first process [-1, 0J as NAAS, while completely removing the annealing reference in the second stage. As we proved in Corollary 3.2, the optimal process of this two-staged AS baseline satisfies the desired target. As shown in Table 7, the two-stage AS (first row) performs similarly to the original AS. In contrast, our NAAS (second row) outperforms this baseline by a significant margin. These results highlight the important role of the annealed reference in NAAS, which contributes non-trivially to its improved performance.
Table 7: Comparison between two-staged AS (Havens et al., 2025) and NAAS. We report Sinkhorn distances (L).
Method
Interpolation
MW54
GMM40
two-staged AS NAAS
U1e(0.1) Uo
0.43
19776.91
U. = (1 U)U + at
0.10
496.48
Effect of a; The diffusion term o, on time E [0, 1] is a key hyperparameter of our method. This controls the guidance scale of annealed reference dynamics, so if o, is large, then gives more guidance, small then small guidance. As shown in Table 8, we ablate the performance of NAAS on GMM40 benchmark by gradually increasing the Omax that defines o+ == " 2 1-t
22ogg dass
max
drinn
As shown in Table 8, it is evident that the performance of NAAS depends on Fmax. as the hyperpa- rameter directly affects its exploration capability. However, once max surpasses some threshold ( 20 for this specific benchmark), the performance of NAAS stabilizes and becomes insensitive to further increase of maax- That said, tuning o; in practice could be straightforwardly performed with e.g., grid or binary search, without having access to samples from or near target distribution. We emphasize these results as strong evidence of the robustness NAAS for practical usages.
Table 8: Ablation studies on mmax- We report Sinkhorn distances (L).
Method mma
max
10
max
20
0 max
50
T maax
100
Tmax
200
NAAS
22893.72
1350.76
639.86
496.48
443.55
492.83
30

---

t- t=0.00
t-0.33
t-0.66
t-1.00
Target
-5
-5
-5
-5
-10
-5
10
-10
-5
10
-10
-10
-5
10
-10
-5
19
(3)
Omaxx
3
1-0.00
F-0.33
Iw0.66
7- 1=0000
Target
-5
-5
-5
10
-5
10
10
-5
10
10
10
-5
10
(b) mmax
20
t-0.00
L-0.3
(=0.66 06.66
(=1.00
Targel
-10
10
(c) mmn
50
Figure 8: Ablation on max- As df increases, the prior distribution μ t becomes increasingly similar to the stationary distribution x e Uo
Visualizing Samples for Various max We visualize the effect of varying the guidance strength, parameterized by σt. As shown in eq. (13), increasing a, drives the prior distribution at O closer the reference distribution, i.e., Uo_ As illustrated in Figure 8, larger values of a, make increasingly resemble a Gaussian distribution N(0, 327) x e -Uo, which is consistent with our theoretical expectation.
Mode Coverage We further conduct experiments to assess whether the proposed sampler can accurately recover mode weights in multimodal distributions with non-uniform component weights. Following the setup introduced in LRDS (Noble et al., 2025), we consider a bimodal Gaussian mixture with components located at (--al,, 21) and (( E,), associated with mixture weights W1 and w2 = 1. respectively.
As reported in Table 9, method achieves competitive performance and accurately recovers the true mode proportions, indicating that NAAS does not exhibit mode blindness or mode switching in this setting. Interestingly, observe that the controlled dynamics in the first stage play a crucial role reweighting the modes. At the beginning of training, the estimated weights approximately vi 0.5, but the control for / [-1, is learned, the sampler gradually reallocates mass toward ú, R 2 aligning with the true distribution.
[-1, 0)
3.
Table 10 summarizes our results across varying between-men mode distances NAAS consistently recovers the correct mode weights and performs comparably to LRDS, even as the separation between modes increases. For LRDS, we retained the default hyprparrameters but trained longer for 1, 3 to ensure convergence; we note that its performance could likely be further improved with additional tuning. We emphasize, however, that these results not intended as extensive comparison between NAAS and LRDS, but rather as an ablation study demonstrating the robustness NAAS in challenging multimodal scenarios with widely separated modes.
We also evaluate NAAS and LRDS on simpler benchmarks, MW54 and Funnel. As shown in Table 11, NAAS achieves slightly better overall performance than LRDS under these settings.
31

---

Table 9:
Comparison mode coverage. We re- port the absolute deviation in mode weight, Vō) wi (1).
Table 10: Ablation study on mode dis- lance with dimension d 16. We report the absolute deviation in mode weight, [01
Table 11: Comparison on MW54 and Funnel. We report Sinkhorn distance (4).
w1 (4).
Method d=
32
Method a= 1
a 2
Method
MW54
Funnel
20.5
LRDS
1.7 %
2.7
LRDS
1.7 %
6.5
6.7 %
LRDS NAAS
0.85
152.09
NAAS
3.3 %
2.0 %
NAAS
1.7
3.2 3.%
3.6 %
132.30
D Implementation Details
In this section, provide the detailed setup for synthetic energy experiments in Table 1. We evaluate our model on four synthetic energy functions, namely MW-54, Funnel, GMM40, MoS following SCLD (Chen et al., 2024). These examples typical and widely used examples to examine the quality of the sampler. Throughout the discussion, let d be a dimension of the state and let
x (I1, 22,
MW-54 We use the same mall-well energy function as in SCLD (Chen et al., 2024). We consider many-well energy function, where the corresponding unnormalized density is defined follows:
v(x) x exp
Σ 1=1
(x - 5)2
(655)
2 m++1
where d = 5, m = 5, and 5 = 4, leading to 2 = 2 = 32 wells in total.
Funnel We follow the implementation introduced in Neal (2003); Chen et al. (2024), which is funnel-shaped distribution with 10. Precisely, defined as follows:
v(x) x N((1;0,, (2) W(22 Ia): exp(xx))))
(66)
with a² 9.
GMM40 Let the target distribution is defined as follows:
v(*) m
(x))
(67)
Following (Blessing et al., 2024; Chen et al., 2024), each mode of Gaussian Mixture Model (GMM) defined as follows:
1/2 = -Vmm,, I), my ~ Ua(-40,
(68)
where Ua(a, b) a uniform distribution on [a, bjd. We use d = 50.
MoS We also use the same Mixture of Student's t-distribution (MoS) following Blessing et al. (2024); Chen et al. (2024). Following the notation of eq. (67), we define MoS of m 10 and d 50, where each mode v4 is defined as follows:
144 12 + m,, my (-1,00 10),
(69)
where to a Student's t-distribution with a degree freedom of 2.
Alanine Dipeptide (AD) molecule that consists of 22 atoms in 3D, therefore, in total 66- dimension data. Following prior work (Zhang and Chen, 2022; Wu et al., 2020). we aim to sample from its Boltzmann distribution at temperature 300K. Moreover, we convert energy function internal coordinate with the dimension d 60 through the OpenMM library Eastman al. (2017).
32

---

Table 12: Hyperparameter Settings
Hyperparameter
MW54
Funnel
GMM40
MoS
AD
K
10
5
5
10
(Nuy N.)
(100, 100)
(100, 100)
(100, 100)
(100, 100)
(1000,200)
(Mus My)
(400, 400)
(400, 400)
(400, 400)
(400, 400)
(2000.2000)
B
512
512
512
512
512
[
10000
10000
10000
10000
10000
N
2048
512
512
512
200
Emnu
100
1000
1000
1000
1000
Amax
100
100
10
(Ir I(p)
(10-6, 10-8)
(10
10
(10-6,
10-00)
(10-1,
10-6)
(10-1, 10-6)
0
1.0
1.0
50
15
1
1.0
9
50
1000
2
man
min
0.01
0.01
0.01
0.01
0.01
gradient clip
1.0
1.0
1.0
1.0
1.0
Training Hyperparameters We adopt the notation of the hyperparameter of Algorithm For additional hyperparameters, let B be the batch size, Ir, and Ir, be the learning rate for u" and v². respectively. Morcover, note that we use Adam optimizer (Kingma and Ba, 2014) with (3, (0, 0.9) for all experiments. Let [B]
[B.. |B,|| be the buffer size for both buffers, K be total number stages (see line 1), and let M,, and M, be a total number of iterations to optimize lines and 12, respectively. Let N be the number of samples generated in lines 3 and 9, hence, total of N samples are eventually added to the buffer in line and 12. Moreover, for stability, we clip the maximum gradient of energy function VU, with clip parameter of Emax and also clip eq. (23) with hyperparameter of Amax- We use geometric noise scheduling from Song et al. (2021); Karras et al. (2022) schedule the noise levels {or [e), controlled by hyperparameters maxx and min (See eq. (64)). The hyperparameters for each experiment are organized in Table 12.
Comparisons and Evaluation metric follow the implementation of (Blessing et al., 2024). The benchmark results are taken from Blessing al. (2024), Chen al. (2024) and Akhound- Sadegh KL divergence (Dxl) and Wasserstein distance between energy values (E(-)W2) of generated and reference samples. For W2 estimation, we use Python OT (POT) library (Flamary et al., 2021). For all metrics, we use 2000 generated and reference samples.
For both the MMD and Sinkhorn distance in Table
(2024). For metrics in Table 3, we evaluate
33


Reviews:
[{"review_id": "Tbhlshn75v", "reviewer": "Reviewer_sTeb", "summary": "his paper presents a Non-Equilibrium Annealed Adjoint Sampler, a stochastic optimal control-based diffusion method that enables scalable training without resorting to importance sampling. To this end, the authors make the following contributions:\n1. The authors introduce a stochastic optimal control-based unbiased sampler where the reference dynamics is governed by an annealed stochastic differential equation. \n\n2. The authors develop an adjoint matching-based optimization algorithm that updates two control functions in alternation, enabling scalable training while avoiding importance sampling and its associated variance issues.", "questions": "1. Could you provide scalability experiments on synthetic problems of increasing dimension, comparing, for example, wall-clock time or memory usage against importance sampling and other baselines?\n\n2. Could you describe whether a joint optimization approach for both control functions is feasible, and if so, provide empirical evidence comparing the alternating scheme against joint optimization? Additionally, could you include a discussion or empirical evidence about  convergence diagnostics and discuss potential failure modes of the alternating optimization scheme?\n\n3. Could you provide empirical variance analysis comparing your method's variance reduction against, for example, importance sampling-based methods?\n\n4. Could you include the missing reference for the SMC baseline used in your experiments and ensure all baselines are properly cited?\n\n5. The paper would benefit from a presentation review to improve clarity.\n\n6. Could you discuss the performance (in particular, generalization) of the proposed method on perhaps more diverse molecular systems beyond the alanine dipeptide?", "limitations": "Yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The paper presents effectively the stochastic optimal control formulation and the adjoint matching methodology.\n\n2. The paper provides strong theoretical results: Theorem 3.1 demonstrates that the SOC-based approach yields the desired sampler, while Theorem 3.4 establishes the crucial unbiasedness property of the sampler.\n\n3. The authors choose appropriate test problems: synthetic multi-modal distributions for rigorous quantitative evaluation and alanine dipeptide generation for practical molecular applications.\n\n4. In the appendix, the authors study the effect of annealing schedule on training performance. \n\nWeakness:\n1. Having two learnable control functions effectively doubles the parameter count, which may limit scalability to high-dimensional problems.\n\n2. While the replay buffer may reduce computational complexity, the implementation of the proposed Algorithm 1 appears complex and could introduce stability issues. The alternating optimization scheme may face convergence challenges that are not discussed.\n\n3. Despite claiming scalability as a key contribution, the paper lacks empirical scalability analysis. No experiments test performance on high-dimensional problems.\n\n4. The paper does not discuss potential failure modes of the alternating update scheme or provide empirical evidence comparing scalability against importance sampling baselines. Additionally, no variance reduction measurements are provided, despite this being a claim in the abstract over importance sampling methods.", "was_revised": true, "final_justification": "The authors have provided a good rebuttal. The responses provide new and valuable experimental data, and agree to make several important clarifications in the manuscript. In particular, the new scalability experiments strengthen their paper. However, the authors are somewhat evasive on key requests for direct empirical comparisons (alternating vs. joint optimization, direct variance measurements), offering theoretical justifications or re-interpreting existing results. While their arguments are plausible, the lack of direct evidence for these specific points remains a minor weakness.", "rebuttal": "We appreciate the reviewer for spending time reading our manuscript carefully and providing valuable feedback. We hope that our responses are helpful in addressing the reviewer’s concerns. Due to the strict character limit, we have grouped related comments and questions where appropriate. \n\n$ $\n\n---\n**1. Clarification on alternating vs joint optimization (W1, W2, Q2)**\n\n*Summary of reviewer’s comments*:\nThe reviewer raised concerns about scalability due to the doubled parameter count from two learnable control functions (W1) and the complexity and stability of Algorithm 1 (W2). They suggest considering a joint optimization approach and discussing convergence diagnostics and potential failure modes (Q2).\n\n*Author response*: \nWe thank the reviewer for these insightful questions. We agree that the choice between joint vs alternating optimization is a central design decision. Here, we provide further clarification on motivation of alternate optimization. \nThe optimization of NAAS occurs in two stages. The first stage, over $t$=[−1,0], learns a broad, smooth prior distribution $\\mu= e^{- U_0 - V_0}$ (where $U_0$ is typically the energy of Gaussian) that broadly covers the target distribution $\\nu$. That is, rather than capturing all fine-scale structures, the first stage aims to construct a low-complexity distribution with broad support over $\\nu$. The second stage, over [0, 1], then refines the dynamics to guide trajectories from $\\mu$ to $\\nu$. As annealed dynamics already bias toward high-density regions, this stage learns to correct residual mismatches and is easy to learn.\n\nDecoupling these objectives improves stability and training. Each stage solves a simpler, focused subproblem: the first creates a good initialization, and the second incrementally corrects the flow. This separation reduces interference, enhancing convergence and sample quality.\n\nFinally, regarding the reviewer’s concern about implementation complexity, we would like to clarify that our algorithmic structure is straightforward. For both 1st and 2nd stages, we (i) roll out trajectories, (ii) compute adjoint variables via standard backward integration, and (iii) update the networks using buffered samples with the stable adjoint-matching loss. These procedures are modular and computationally efficient. \n\nWe have incorporated a discussion of these design choices, results, and interpretation in the Appendix and have also referenced it in the main text for clarity.\n\n$ $\n\n---\n**2. Additional experiments on higher-dim problems (W3, W4, Q1)**\n\n*Summary of reviewer’s comments*:\nThe reviewer noted a lack of empirical scalability analysis (W3), requested comparison on higher-dim problems with other baselines (Q1), and potential failure modes (W4).\n\n*Author response*: \nWe thank the reviewer for highlighting the importance of evaluating scalability and for suggesting a more detailed empirical analysis on high-dimensional settings. \n\nWhile we agree that scalability is a crucial aspect, we kindly note that our current experiments readily include challenging, high-dimensional problems in sampling literature [1-4] where ground-truth samples are obtainable. Specifically, the GMM40 and MoS datasets in Table 2 are 50 dim, whereas the alanine dipeptide in Table 3 has 60 dim. In all settings, our NAAS consistently outperforms existing baselines by a substantial margin, achieving over 2× improvement in sample quality over the second-best method. Importantly, we compare against a wide range of importance sampling–based baselines, including SMC, SMC-ESS, and SCLD [4]. Our method not only achieves better average performance and more robust convergence across random seeds.\n\nIn response to the reviewer’s suggestion, we conduct additional experiments on 100-dim GMM40 tasks. Specifically, we evaluated our method and two strong baselines: Adjoint Sampling [9], which uses standard reference dynamics, and SCLD [4], a SOTA importance sampling method. We also report wall-clock time and peak memory usage for 50-dim experiment, measured on 1 A6000 GPU. The results are summarized below:\n\n| Method | $d=4$ | $d=16$ | $d=50$ | $d=100$ | Mem. usage. | run time |\n|----|----|----|----|----|----|----|\n| SCLD | 157.90 | 1033.19 | 3787.73 | 12357.49 | - | 30-40 mins |\n| AS | -  | - | 18984.21 | - | 1GB | 30 mins |\n| Ours | **25.15** | **37.96** | **496.48**  | **633.20** | 1GB | 120 mins |\n\nOur method remains robust and maintains strong performance even as dimensionality increases, while the baseline, SCLD, experience significant degradation as the dimensionality grows. Furthermore, our method demonstrates competitive runtime and memory consumption, remaining scalable as dimensionality increases. This is partly due to the use of replay buffers, which effectively ensures efficient training. We acknowledge, however, that a limitation of our method is the need to solve a backward adjoint state (Equation 22), which contributes to the longer training time. Nonetheless, we believe that the resulting performance gains, especially in sample quality, more than justify this additional computational cost.\n\nRegarding potential failure modes of the alternating scheme, we would like to note that across all tasks, including the high-dimensional ones, we did not observe signs of divergence or instability when properly tuned. We attribute this to the structured separation of roles in the alternating updates, which helps stabilize training dynamics, as discussed in our response to Q2. \n\n$ $\n\n---\n**3. Discussion on variance reduction (W4, Q3)**\n\n*Author response*: \nWe thank the reviewer for raising this important point. We recognize that the phrasing in the abstract and introduction “reliance on importance sampling leads to high variance” may have misled readers. Hence, we would like to clarify on the high-variance property in importance sampling (IS)–based algorithms and why our method is less susceptible to this issue.\n\nFormally, IS methods approximate the target distribution by $\\nu\\approx\\sum^N_{i=1} w^{(i)}\\delta_{X^{(i)}}$, where the unnormalized weights are obtained by $w^{(i)} \\propto \\frac{dp^\\star}{dp} (X^{(i)})$ with the ratio between the target path measure $p^\\star$ and proposal $p$. When the support or mass of the proposal and target distributions are misaligned—as is common in high-dimensional or complex energy landscapes—these weights can suffer from high variance, as noted in prior works (e.g. [5]; [6] Section 4.1.4). This leads to poor effective sample size and significant variance in the empirical approximation, and can be further exacerbated when the number of samples $N$ is insufficient.\n\nFor example, in PDDS [3], the learned distribution approximates a reweighted empirical measure and is thus exposed to the instability induced by high-variance weights. In contrast, our method learns a sampling distribution directly via path-space control, thereby avoiding the need to compute or rely on importance weights.\n\nAdditionally, as shown in Table 2, SMC-based methods (including SMC, CRAFT [7], and SCLD [4]) exhibit substantially higher variance on the 50-dimensional GMM40 and MoS benchmarks, especially in comparison to PIS [8], AS [9], and our method, which do not rely on importance weighting. We will clarify this point in the revised manuscript and provide additional discussion to explicitly contrast variance behaviors.\n\n$ $\n\n---\n**4. Clarification on remaining questions**\n\n**Q4**: *Missing SMC references*\n\n*Author response*: We thank the reviewer for pointing this out. We follow the SMC and SMC-ESS implementation in [4]. We will incorporate the explanation in Line 238.\n\n**Q5**: *Presentation*\n\n*Author response*: We thank the reviewer for their detailed comments, including the helpful suggestion regarding presentation clarity. Following the reviewer’s feedback, we have revised the manuscript to improve overall clarity and structure. In particular, we have \nclarified the stability and scalability aspects of our method,\nExpanded the discussion on design choices (e.g., alternating vs. joint optimization), and\nIncorporated additional discussion on importance sampling and its variance\nas raised in W1-5. We again thank the reviewer for helping us clarify and improve the manuscript. We believe these revisions significantly enhance the clarity and accessibility of our work, and better highlights our key contributions.\n\n**Q6**: *Generalization of NAAS beyond alanine dipeptide (ALDP)*\n*Author response*: We agree with the reviewer that evaluating the generalization of our method on more diverse and larger-scale molecular systems beyond ALDP would further strengthen the paper. Due to the limited time available during the rebuttal period and the significant computational cost associated with high-dimensional molecular simulations, we were not able to include such additional experiments. However, we acknowledge this as a limitation and have added a corresponding discussion in the revised manuscript. \n\nOn the other hand, we would like to note that ALDP remains a widely used and nontrivial benchmark in the molecular sampling literature. With approximately 60 dimensions, it already presents meaningful challenges.\n\n---\n**References**\n\n[1] Improved Sampling via Learned Diffusions\n\n[2] Non-equilibrium Transport Sampler\n\n[3] Particle Denoising Diffusion Sampler\n\n[4] Sequential Controlled Langevin Diffusions\n\n[5] Sequential Monte Carlo Samplers\n\n[6] Free Energy Computations: A Mathematical Perspective\n\n[7] Continual Repeated Annealed Flow Transport Monte Carlo\n\n[8] Path Integral Sampler\n\n[9] Adjoint Sampling"}, {"review_id": "yTXn61KU2T", "reviewer": "Reviewer_YsLF", "summary": "In this work, the authors present a novel Adjoint Matching approach for sampling from a given unnormalised density. This aligns closely with recent interest in the generative modelling community in moving away from sample-based approaches. Within this framework, the presented theory falls under the category of using a controlled annealed SDE. However, in contrast to most work that uses the annealing interpolation for reweighting, the presented method proposes an interesting approach to provide unbiased samples by considering how to transform the initial distribution.\n\nThe proposed method achieves the learning of the controlling force for the SDE, as well as how to sample from the adjusted prior, using two different variants of the Adjoint Matching framework. These two sides of the learning problem are then combined in alternating fashion. A replay buffer is used to reduce simulation overhead.\n\nThe method is evaluated on a number of well-established benchmarks, on which it demonstrates excellent performance.", "questions": "Please refer to the Strengths and Weaknesses section for a complete discussion of my comments and questions. As stated there: *Pending the authors’ response to my comments on the motivation of the work and the comparisons with reweighting taken into account, I do not feel confident recommending the paper for publication at this time. However, if these points are properly addressed, I would be willing to increase my score to a 5. I will likely not increase my score to a 6, as the work builds heavily on existing variants of adjoint matching and contains limited real-world evaluation. Please see below for a more detailed list of comments and questions.*", "limitations": "As noted by the authors in the conclusion, the method still needs to be tested further on real-world problems. I do not foresee any negative societal impact at this time.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper presents an interesting approach to an important and currently highly active field of research: sampling from a distribution given its unnormalised density. The paper is generally well written, contains a great deal of technical depth, and demonstrates strong results on a number of well-established benchmarks.\n\nHowever, the motivation for the proposed approach, and with it the main contribution, is not clearly communicated in the writing. My understanding is the following. Given an SDE with time-dependent density $U_t$, recent work has focused on using importance weighting to debias the drawn samples. This, however, can result in high variance. Adding a control term $u$ can help reduce this variance, but it also requires modifying how the weights are calculated. Specifically, it introduces an additional divergence term and a cross term. These modifications can reduce the variance of the weights, but unless $U_t$ and $u$ together satisfy the continuity equation, the weights will still exhibit some variance. Related work addresses this by directly learning to minimise the controlled weights, and with that the variance of the estimator. In the context of Adjoint Matching, however, minimising the controlled weights is not part of the learning objective. As a result, even after convergence, the weights are not constant. One way to resolve this is to adjust the prior distribution, which is the focus of this paper.\n\nAs a side note, and this could be more clearly stressed in the paper, even in the case where $U_t$ and $u_t$ are learned to solve the continuity equation, there remains a $Z_0 / Z_1$ term. The only way to eliminate this term is to adjust the prior. Therefore, changing the prior is not only a way to avoid varying reweighting factors, it is a necessary step to remove the constant weight factor altogether.\n\nIf my understanding is correct, I believe this should be explained more clearly in the paper. With this change, the motivation and contributions would become more transparent. The current primary motivation of avoiding reweighting is, in my opinion, not convincing. If calculating the weights is feasible, this always offers an advantage. Given two samplers that provide similar sample quality, a sampler that supports reweighting is always more useful than one that does not.\n\nThis also relates to another aspect of the paper that I believe should be reconsidered. In the experiments section, the paper compares several baseline methods, including both those that support reweighting and those that do not. However, the reweighting is not taken into account in any of the evaluations. This is particularly notable for SCLD, the second-best performing method, which does support reweighting. I would be interested to see a comparison on a task where reweighting is beneficial, for example, estimating an expectation.\n\nPending the authors’ response to my comments on the motivation of the work and the comparisons with reweighting taken into account, I do not feel confident recommending the paper for publication at this time. However, if these points are properly addressed, I would be willing to increase my score to a 5. I will likely not increase my score to a 6, as the work builds heavily on existing variants of adjoint matching and contains limited real-world evaluation. Please see below for a more detailed list of comments and questions.\n\n**Further comments and questions:**\n\n- The paper (and so does prior work) frequently uses the notion of  \"non-equilibrium,\" including in the title. I would appreciate a clearer discussion of what criteria make the method non-equilibrium, and whether these still apply to the proposed approach. Does \"non-equilibrium\" refer to intermediate instantaneous distributions not matching the distribution defined by $U_t$? The concept of equilibrium becomes more confusing when non-conservative forces such as $u$ are present, as they affect the stationary distribution of the system.\n\n- **Table 1:** The use of check marks and large red crosses, along with the phrasing \"reliance,\" implies a strong negative stance on the use of importance weighting. I do not believe this is justified. Changing this to a simple \"Uses IWS\" with a yes/no column would, in my opinion, be more balanced.\n\n  **L52:** The statement \"ensures unbiasedness by correcting the prior distribution\" is surprising here, as there has not yet been any discussion of how adjusting the prior contributes to unbiasedness. A little bit or rewriting could help improve the readability of the introduction significantly here.\n\n- The notation used to index functions by time is not consistent. Most of the time a subscript is used, but there are also cases where time is passed directly to a function. If this was a deliberate choice, it should be clarified.\n\n- **L110–111:** It is not immediately obvious why the value function at initial time becomes constant. A short clarifying remark would help the reader.\n\n- **L156:** The phrase \"special interpolation\" is unclear, especially since the schedule used is not actually an interpolation.\n\n- **Corollary 3.2:** The claim that DDS is a special case of the proposed method seems questionable. DDS does not use an interpolation and does not alter the prior.\n\n- **Figure 1:** This figure was very helpful. Including a similar illustration in the introduction could greatly improve clarity and reduce early confusion.\n\n- **L207–209:** The discussion of the replay buffer appears to repeat what was said in the previous paragraph.\n\n- **Table 2:** A minor suggestion, but linking each method to its citation would improve the table’s usefulness.\n\n- Regarding the baselines, I would expect NETS to be included given the strong methodological overlap. If it was omitted due to unavailable code, this should be explicitly stated.\n\n- Related, it would be useful for the authors to explain why the MMD scores are not reported for most methods on the Funnel and MoS datasets.\n\n- Many in-text citations do not correctly link to the references at the end of the paper. Clicking them takes the reader to the first or second page instead.", "was_revised": true, "final_justification": "In their rebuttal and following discussions the authors clarified a number of points of their paper and with that further established their contribution in my opinion. Based on the initial rebuttal I've increased my score to a 4, and consecutively increased it to a 5 after their final set of comments.", "rebuttal": "We appreciate the reviewer for spending time reading our manuscript carefully and providing valuable feedback. We hope that our responses are helpful in addressing the reviewer’s concerns. Due to the strict character limit, we have grouped and summarized related comments where appropriate.\n\n**1. Clarification on motivation**\n\n*Summary of reviewer’s comments*:\nThe reviewer raised the unclear motivation in the current presentation and hypothesized that the motivation may be to propose a diffusion sampler (e.g., NETS [1]) that learns a control $u$ satisfying continuity equation prescribed by the annealing energy $U_t$. Since the proposed method does not minimize variance of weights explicitly (as like LogVar [2]), additional adjustment on prior distribution is necessary. The reviewer also raised questions about reweighting.\n\n**A.** We thank the reviewer for the effort to understand our manuscript. We agree that our motivations need improvement. In this response, we address points that may have been misleading or unclear (1.1-1.3). Based on this clarification, we have revised the corresponding parts of the manuscript to improve transparency and presentation (1.4).\n\n**1.1 The roles of $u$ and $U_t$**\n\nThe reviewer understands correctly that, recently, there have been a few annealed-based diffusion samplers that learn a control $u$ to enforce the continuity equation by introducing reweighting terms to align the marginals. A noticeable example is NETS [1]. That is, the role of $u$ for methods like NETS is to learn a control $u$ that matches prescribed marginal $p_t \\propto e^{-U_t}$ given annealing energy $U_t$.\n\nHowever, the aforementioned motivation differs significantly from our proposed method, NAAS. Specifically, the optimal control considered in NAAS needs not match $p_t \\propto e^{-U_t}$, except at the terminal time t=1 (see Theorem 3.4). The distinction stems from the fact that, unlike NETS which does not admit Stochastic Optimal Control (SOC) interpretation, NAAS is grounded fundamentally on the SOC framework whose control $u$ is directly optimized to reach the final target distribution, without being constrained to follow any specific path of intermediate marginals. Moreover, thanks to the rich literature on solving SOC problems, we have the flexibility to choose methods, including Adjoint Matching (AM), that do not rely on importance sampling.\n\n**1.2 Discussion on weight-based vs AM-based SOC solvers**\n\nIn the case where SOC problems are considered for Boltzmann sampling, there exist prior methods (such as LogVar [2]) that solve the SOC by directly minimizing the variance of the importance weights induced by the controlled dynamics, as mentioned by the reviewer. In contrast, other approaches—including our NAAS—use AM [3,4], which optimizes a matching loss based on the backward adjoint state.\n\nWe emphasize that both solvers aim to approximate the same optimal control. Therefore, if both methods are solved optimally, they yield identical solutions, hence, importance weights of optimal solution should be constant. In our work, we adopt AM due to its empirical advantages in variance reduction, scalability, and learning stability.\n\nWe also emphasize that employing an AM-based SOC solver does not exclude the use of resampling during inference. We agree with the reviewer that resampling can be particularly beneficial in the evaluation phase, especially when samples produce low-variance estimators, as discussed in Appendix B. While resampling could theoretically be incorporated during NAAS training, in practice, we find that using AM alone is sufficient to achieve strong performance.\n\n**1.3 Why do we need to learn the prior distribution**\n\nThe SOC problem in Theorem 3.2 prescribes an initial distribution $\\mu$ (see Equation (13)) such that the corresponding terminal distribution matches the target $\\nu$. However, $\\mu$ is generally intractable. To address this, we introduce an auxiliary SOC problem (Lemma 3.3) to learn a prior distribution that approximates $\\mu$. Then, in Theorem 3.4, we show that these two SOC problems can be unified into a single joint problem whose solution produces samples from the target distribution $\\nu$ when optimally solved.\n\n**1.4 Clarification of our motivation & contribution**\n\nRather than positioning our method solely as an annealing-based sampler that avoids importance sampling, we will present it as one of the key characteristics of our approach. Instead, we would like to clarify that the core motivation and contribution of our work lies in developing a SOC-based sampler that incorporates annealed reference dynamics, as illustrated in the following flow:\n\n- Prior SOC-based approaches [4,7,8] have constructed unbiased diffusion samplers using stationary or uninformative base references, and do not fully leverage annealing references as employed in other importance sampling based methods. In this work, we propose a novel SOC-based framework that employs annealed reference dynamics as the base SDE.\n- (*) Precisely, we formulate two SOC subproblems: (i) learning a prior distribution $\\mu$ from the initial distribution $\\delta_0$ with standard base reference, and (ii) transporting $\\mu$ to the target $\\nu$ using controlled annealed dynamics. We unify these into a single SOC formulation, which yields a controlled diffusion process that effectively transports $\\delta_0 \\to \\mu \\to \\nu$.\n- Among available solvers for SOC, we adopt AM [3], a recently proposed method being scalable and sample-efficient. As discussed in Adjoint Sampling [4], this approach enables scalability and improved training stability. \n\n**2. Clarification on reweighting**\n\n**A**. Thank you for raising this important point. We realize that our manuscript may have caused some confusion regarding the role of reweighting. Specifically, when we stated that a method \"relies on importance sampling,\" we were referring to the use of importance sampling during training. Our key claim is that, our method does not require importance sampling during training. This distinction is important because using importance sampling during training often leads to high-variance [5,6], which can negatively affect performance. We will revise the manuscript to clarify this distinction in the abstract, introduction, and Table 1.\n\nRegarding your point about evaluation: we appreciate the suggestion to consider reweighting at test time. Like other SOC-based samplers [4,7], our method does support reweighting during evaluation, as discussed in Appendix B. However, in this work, we chose not to emphasize evaluation-time reweighting to focus on the properties of the learned samplers themselves. We agree that incorporating reweighting into evaluation could be informative and plan to explore this direction in future work.\n\n**3. Clarification on remaining questions**\n\n*Q1*. The meaning of “non-equilibrium”.\n\n*A.* In equilibrium-based samplers such as Langevin dynamics, the stochastic process is constructed so that its stationary distribution coincides with the target distribution. Sampling is achieved by simulating the process over a sufficiently long (often asymptotic) time horizon. These dynamics eventually “equilibrate” to the target, and accurate sampling is guaranteed only in the infinite-time limit. In contrast, non-equilibrium samplers explicitly construct a finite-time dynamical system whose terminal distribution matches the target, without requiring convergence to a stationary state. In our method, NAAS, we formulate the sampling task as a SOC problem, where the control is optimized to ensure that the terminal distribution at time $t=1$ matches the desired target distribution $\\nu$. \n\n*Q2*. In Table 1, phrasing \"reliance,\" implies a strong negative stance. Changing this to a simple \"Uses IWS\" with a yes/no column would be more balanced.\n\n*A*. We agree with the reviewer’s comment. We revised the table to use the more neutral label “Uses IWS during training”.\n\n*Q3*. L52: \"ensures unbiasedness by correcting the prior distribution\" is surprising here, as there has not yet been any discussion of how adjusting the prior contributes to unbiasedness. A little bit or rewriting could help improve the readability.\n\n*A*. We agree that introducing the notion of \"unbiasedness\" at that point in the introduction was insufficiently motivated. To enhance clarity and improve the logical flow, we revised the corresponding paragraph, now presented as paragraph (*) in the “1.4 Clarification of Our Motivation & Contribution” section.\n\n*Q4*. The notation used to index functions by time is not consistent.\n\n*A*. As the reviewer suggested, we put every time variable into the subscript. For example, we fixed time variables in adjoint terms and neural networks as subscripts.\n\n*Q5*. Typo / Minor comments on L110–111, L156, Corollary 3.2, Figure 1, Table 2, citation/reference.\n\n*A*. We thank the reviewer for carefully reading our paper and providing valuable suggestions. We will reflect all these suggestions to clarify our manuscript. Due to the character constraints of this rebuttal, we are unable to detail all the revisions here. However, we would like to highlight that most of the metric values in Table 2 were taken from the existing papers, and MMD scores are not reported for many methods in their original works. We will include this explanation in our manuscript to ensure clarity. Additionally, we will correct the link to the reference as suggested.\n\nOnce again, thank you for your careful review and thoughtful suggestions. We hope that the revisions we’ve made address your concerns.\n\n---\n[1] Non-equilibrium Transport Sampler\n\n[2] Solving high-dim HJB PDEs using neural networks\n\n[3] Adjoint Matching\n\n[4] Adjoint Sampling\n\n[5] Sequential Monte Carlo Samplers\n\n[6] Free energy computations: A mathematical perspective\n\n[7] Path Integral Sampler: a stochastic control approach for sampling\n\n[8] Denoising Diffusion Samplers\n\n[9] Sequential Controlled Langevin Diffusions"}, {"review_id": "ye1HSDHKBe", "reviewer": "Reviewer_GpP4", "summary": "This paper proposes a novel variational diffusion-based sampler defined on time interval $[0,1]$, that aims to sample at $t=1$ from an unnormalized target density $\\pi$ (without access to ground truth samples), by leveraging \"reference\" stochastic dynamics informed by the target distribution. The goal is to learn an additive control to the \"reference\" drift, to steer the particles towards the high density areas (which are already given by the reference dynamics) but with the right proportions (which is the harder task). In line with recent works [1, 2], the authors consider the Annealed Langevin dynamics as reference, where the time-evolving density $p_t^{\\text{ref}}$ is a geometric interpolation between a standard Gaussian distribution and $\\pi$. However, the formulation of the problem is different : while previous works meant at correcting the reference dynamics to marginally match $p_t^{\\text{ref}}$ via a variational formulation, the current approach relies on a novel Stochastic Optimal Control (SOC) formulation which enables to exactly match $\\pi$ at time $t=1$ but not $p_t^{\\text{ref}}$ for $t<1$. In particular, the induced SOC loss relies on a state cost and not a final cost, as opposed to previous (memoryless) SOC-based diffusion samplers, and can be optimized via the backward Adjoint Matching procedure [3]. In this specific case however, the starting marginal at time $t=0$ is intractable and needs to be approximated to obtain starting samples. To circumvent this issue, the authors propose to extend the SOC to $[-1,0]$, by formulating a standard memoryless setting (ie denoising diffusion), which can be solved via the Adjoint Sampling method [4]. In the end, the overall algorithmic procedure alternates between minimizing these two objectives (hence learning two different neural networks), which can be made by sharing the use of buffers. This method is compared with annealed MCMC methods and recent diffusion-based samplers on toy mixture settings and molecular systems.\n\n[1] Transport meets Variational Inference: Controlled Monte Carlo Diffusions. Vargas et al, 2024\n\n[2] Sequential Controlled Langevin Diffusions. Chen et al. 2025\n\n[3] Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control. Domingo-Enrich et al. 2025\n\n[4] Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching. Havens et al 2025.", "questions": "**Questions**\n- Following my remark in the Section *Weaknesses* on Figure A, where the trickiest part of the multimodal problem seems to be solved in the first interval $[-1,0]$, how can you handle in practice the tradeoff of \"hardness\" (i.e. controlling carefully the mode speciation) between the two SOC parts ? Does it relate to the setting of $\\sigma_t$ ? Can it rely on realistic hyperparameter tuning (ie without access to ground truth samples) ?\n- How do you initialize the buffers in practice ? This is often the tricky part in diffusion samplers. For instance, LRDS leverages the existence of Monte Carlo samples collected in the target modes (but not informed by the relative weights) to initialize the reference drift; AS also uses this technique to evaluate for the first time its variational loss. These approaches hence ensure fast convergence. In the proposed setting, I do not see how such target-informed samples could be used due to the fact that we have no information on the marginal distribution at time $t=0$, which may lead to slow convergence. Do I miss something ?\n\n**Suggestions**\n- one of the main motivations of this work is to avoid the use of importance sampling schemes, which seems to be denoted as something negative. Although I understand that indeed, IS can lead to high variance when the target and the proposal do not match, even more when the dimension is large; however, resampling is not necessarily doomed to fail in general settings, if the target and the proposal are well designed. Most of real-world samplers rely on versions of Parallel Tempering, where the resampling steps work because the intermediate temperatures are carefully chosen to avoid huge mismatch between densities. Moreover, resampling may help to recalibrate the diffusion model with respect to the target distribution, hence avoiding accumulating too much numerical bias when generating from the neural sampler (especially at the very end, where we want to be sure to sample from $\\pi$). I think that this motivation should be more nuanced.\n- I recommend the authors to consider multi-modal targets with non-uniform mode weights where the measurement of mode weight recovery (which is non-trivial at all for diffusion-based samplers) can be explicitly computed, see for example [1]. This would enable to understand if the proposed sampler suffers from mode switching and/or mode blindness issues (as mode discovery is already partially handled by the reference dynamics).\n- the benchmark should include (i) LRDS, which is the natural competitor of the current approach for a memoryless setting and (ii) AS, to understand the improvements of using the reference Annealed Langevin dynamics in the second SOC part (more informally, can AS already do the job ?)\n- following my last point, I think it would be interesting to have a visual inspection of the marginal distribution at time $t=0$ depending on several hyperparameter settings, to understand if it already captures the multimodality features of the target, and therefore if the use of annealed reference dynamics is that useful\n\n[1] Improving the evaluation of samplers on multi-modal targets. Grenioux et al, 2025.", "limitations": "To apply the Adjoint Matching method, one needs to have access to the Hessian of $\\log\\pi$, which may be costly to derive; plus, one needs to train two neural networks in the end. This computational burden should be much more emphasized as a limitation of the current method in a context where the computational budget of diffusion-based samplers is questioned, see [1].\n\n[1] No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers. He et al, 2025.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- the paper is technically solid, the introduction of the losses is well detailed and carefully justified\n- the paper is overall well written, the figures explain well the main ideas\n\n**Weaknesses**\n- the current paper does not cite a close competitor LRDS [1] (which was publicly available more than 6 months ago), which aims at solving the same kind of task (i.e., learning a diffusion based sampler by leveraging reference dynamics that is informative on the target without introducing importance sampling scheme) in the memoryless setting (ie with denoising diffusion), while focusing on multi-modal targets. This competitor should be compared with the current method to weigh the pros/cons between the two approaches (in particular, LRDS only learns one control term): namely, should we consider a reference dynamics given by a denoising or a tempering path ?\n- the current method should also be compared with the Adjoint Sampling method in the numerics (ie the classic memoryless setting, yet with non-informed reference dynamics), to understand how incremental the improvement is by allowing the use of reference dynamics with tempering path (which increases the problem complexity). In particular, the results of Figure 1 might suggest that **the mode speciation** of the target density (ie the trickiest part when sampling from multimodal targets starting from unimodal distributions) **may already happen on $[-1,0]$** (ie when applying AS method), resulting in separated modes in the AM part : in this specific case, the annealed reference dynamics may not be much necessary (applying Langevin may be sufficient to capture the local landscape)...\n- the issue of *mode switching*, which is inherent to annealed sampling, is not discussed : it relates to the fact that the mode weights may vary a lot on a given tempering path (which is not the case for denoising paths), which is one of the motivations stated in [1] ti use a denoising path. This may notably result in obtaining wrong mode proportions (while still good mode locations) by the end of the sampling. This has notably motivated previous work [2] to learn the intermediate temperature levels, to be able to suffer less from mode switching (see their Figure 8). Here, the choice of the toy multimodal targets presented in the numerics does not enable to see if the presented method avoids this phenomenon (which may definitely happen since the chosen geometric interpolation is the linear one and is not optimized). Remark that the sampling results on GMM40 (see Figure 2.a) seem to indicate that the true mode proportions are indeed not well recovered by the algorithm...\n- in the same manner, the issue of the *mode blindness of score-based losses* (here, the AS loss on $[-1, 0]$) is not addressed: in the case of well separated modes, the score of the target density is known to be blind to the mode weights (contrary to the log density), see [3]. Hence, for a multimodal target with non equilibrated modes, this blindness of $\\nabla \\log \\pi$ may prevent to reach the target when using AS methodology and rather end up with obtaining samples from an equilibrated version of $\\pi$ (same mode locations, but wrong mode proportions). The current target settings do not enable to observe if the proposed sampler solves this issue.\n- while the proposed approach seems to rely on careful hyperparameter tuning (namely $\\sigma_t$), which, as far as I understand, may be critical to split the task hardness between the two SOC problems (see the section *Questions*). No indication of setting in realistic sampling frameworks is provided in the main paper, which hurts the applicability to real word problems\n\n[1] Learned Reference based Diffusion Sampler for multi-modal distributions. Noble et al, 2025.\n\n[2] Sequential Controlled Langevin Diffusions. Chen et al. 2025\n\n[3] Towards Healing the Blindness of Score Matching. Zhang et al. 2022.", "was_revised": true, "final_justification": "I have increased my score, as the authors have provided convincing answers that will be added to the paper (see my final answer for more details).", "rebuttal": "We appreciate the reviewer for spending time reading our manuscript carefully and providing valuable feedback. We hope that our responses are helpful in addressing the reviewer’s concerns. Due to the strict character limit, we have grouped and summarized related comments where appropriate.\n\n---\n**S1**. One of the main motivations of this work is to avoid the use of importance sampling schemes.  Although I understand that indeed, IS can lead to high variance when the target and the proposal do not match, even more when the dimension is large; however, resampling is not necessarily doomed to fail in general settings, if the target and the proposal are well designed. Moreover, resampling may help to recalibrate the diffusion model with respect to the target distribution, hence avoiding accumulating too much numerical bias when generating from the neural sampler.\n\n**A**. We thank the reviewer for raising an important point regarding the role of importance sampling in our motivation. We agree that the current presentation lacked sufficient justification and clarity on this aspect. In response, we have revised the introduction to better highlight our core contribution: the development of an SOC-based sampler with annealed reference dynamics, which is a strong base dynamics that facilitates transport toward the target distribution. And among various solver of the SOC-based sampler, we employ Adjoint Matching [4], which is known to be scalable, has low-variance and efficient [3]. Precisely, we will revise the introduction to reflect the following structure:\n\n> Prior SOC-based approaches such as PIS [1], DDS [2], and AS [3] aim to solve the SOC formulation (e.g., Equation (8) or (9)) to construct unbiased diffusion samplers. However, these methods typically learn the controlled dynamics starting from a stationary or uninformative base reference, and therefore do not fully leverage annealing references as employed in other importance sampling based methods. \n\n> Motivated by this limitation, we propose a new SOC-based framework that employs annealed reference dynamics as the base SDE. This annealing structure provides a natural progression toward the target distribution and generates more informative reference trajectories, thereby significantly improving the stability and efficiency of learning the control. As we demonstrate empirically, this design leads to improved sample quality and target matching, especially compared to methods that rely on static references.\n\n> To delve deeper into our method, we formulate two SOC subproblems: (i) learning a prior distribution $\\mu$ from the initial distribution $\\delta_0$ with standard base reference, and (ii) transporting $\\mu$ to the target $\\nu$ using controlled annealed dynamics. We unify these into a single SOC formulation, which yields a controlled diffusion process that effectively transports $\\delta_0 \\to \\mu \\to \\nu$, resulting in a sampler whose terminal distribution matches the target.\n\n> Among available solvers for SOC, we adopt Adjoint Matching [4], a recently proposed method being scalable and sample-efficient. As discussed in Adjoint Sampling [3], this approach enables low-variance and improved training stability. \n\n> To summarize, we introduce a novel SOC-based diffusion sampling framework that leverages annealed reference dynamics to enhance learning of the control. This framework, which we call NAAS (Non-equilibrium Annealed Adjoint Sampler), is the first to formulate an SOC problem over diffusion processes with a non-stationary, annealed base. Our approach combines principled formulation with a practical and scalable solver, and demonstrates strong empirical performance across several challenging sampling tasks.\n\n\n---\n**W1-5, S2**. Consider multi-modal targets with non-uniform mode weights where the measurement of mode weight recovery. This would enable to understand if the proposed sampler suffers from mode switching and/or mode blindness issues. The benchmark should include (i) LRDS (ii) AS, to understand the improvements of using the reference Annealed Langevin dynamics in the second SOC part. \n\n**A**. We thank the reviewer for the thoughtful and constructive suggestions. We fully agree on the importance of evaluating whether the proposed sampler can accurately recover mode weights in multi-modal distributions with non-uniform weights. As suggested, we adopt the experimental setup introduced in LRDS [2], focusing on a two-mode Gaussian mixture model (GMM) where the first mode has a true weight of $w_1=2/3$​. We report the absolute error $|\\hat{w}_1-w_1|$, where $\\hat{w}_1$​ is the empirical weight estimation of the first mode.\n\n| Method ($\\downarrow$) | $d=16$ | $d=32$ |  \n|----|----|----|\n|AS| 20.5\\% | - |\n|LRDS| **1.7**\\% | 2.7\\% |\n|Ours| 3.3\\% | **2.0**\\% |\n\nAs shown in the table, our method demonstrates competitive performance and recovers the correct mode weights with high accuracy. This indicates that our method does not suffer from mode blindness or switching in this scenario. Interestingly, we observed that the controlled dynamics in the first stage (from time $[-1,0]$ plays a key role in reallocating the weight. At the beginning of the training, $\\hat{w}_1 \\approx 0.5$, but as we learn the control for time $[-1,0]$, it gradually reallocates weights towards $\\hat{w}_1 \\approx 2/3$. \n\nWe also note that both AS and our method converged stably with a small learning rate of $10^{-6}$. However, due to the limited rebuttal period, we could not perform an extensive hyperparameter search for AS, which may partially explain its degraded performance. We plan to conduct a more comprehensive hyperparameter search for AS and include the results in the revised version of the manuscript.\n\nTo further verify the benefit of incorporating annealed dynamics, we additionally applied AS to the main benchmark tasks from Table 2. As shown below, our method consistently outperforms AS by a significant margin, further supporting the effectiveness of annealed reference dynamics:\n\n| Sinkhorn | MW54 | GMM40 | MoS|\n|----|----|----|----|\n| AS | 0.32 $\\pm$ 0.06 | 18984.21 $\\pm$ 62.12 | 2178.60 $\\pm$ 54.82 |\n| Ours | **0.10** $\\pm$ 0.0075 | **496.48** $\\pm$ 27.08 | **394.55** $\\pm$ 29.35 |\n\n\n---\n**W6, S3, Q1**. Impact of hyperparameter, especially $\\sigma_t$ (W6). It would be interesting to have a visual inspection of the marginal distribution at time $t=0$ depending on several hyperparameter settings (S3). Following my remark in the Section Weaknesses on Figure A, where the trickiest part of the multimodal problem seems to be solved in the first interval [-1,0], does it relate to the setting of $\\sigma_t$ (Q1)?\n\n**A**. Regarding the visual inspection of the marginal distribution at $t=0$, due to rebuttal policies, we are currently unable to include the visualizations. However, we have verified that increasing the guidance strength (i.e., larger $\\sigma_t$​) causes the distribution $\\mu$ at $t=0$ to resemble the reference distribution$\\propto e^{-U_0}$​, which aligns with equation (13). We will include these visualizations and a more detailed discussion in the revised version of the manuscript.\n\n---\n**W7**. No indication of setting in realistic sampling frameworks is provided in the main paper, which hurts the applicability to real word problems\n\n**A**. We agree that evaluating our method on more realistic, large-scale applications (such as molecular generation) would further enhance its practical applicability. However, we would like to kindly note that our current experiments readily include challenging, high-dimensional problems in sampling literature [3,4,7,8] where ground-truth samples are obtainable. Specifically, we include the GMM40 and MoS datasets in Table 2 which are 50 dim, and the alanine dipeptide in Table 3 has 60 dim.\n\n---\n**Q2.** How do you initialize the buffers in practice? I do not see how such target-informed samples could be used due to the fact that we have no information on the marginal distribution at time $t=0$, which may lead to slow convergence.\n\n**A**. We do not use any special initialization for the buffers. For each epoch, we generate a certain number of trajectories (in most of our experiments, we used 512 trajectories) sampled from the range [-1, 1]. As the reviewer correctly noted, we do not have access to the prior distribution $\\mu$ at time $t = 0$. To address this, we obtain samples by starting from $X_{-1}\\sim\\delta_0$. Please note that adjoint-based methods only requires (on-policy) samples, so there is no need to get correct samples at $t=0$.\n\n---\n**L1** Limitations Several drawbacks: (1) requires access to the Hessian of $U_t$ to apply Adjoint Matching, (2) requires two neural networks (see [6]).\n\n**A**. We agree with the reviewer that we should discuss the limitation that pointed out. We will incorporate it in the conclusion section. \n\n---\nOnce again, we appreciate the reviewer for careful consideration and important suggestions. We hope that our responses are helpful in addressing the reviewer’s concerns.\n\n\n---\n**References**\n\n[1] Improving the evaluation of samplers on multi-modal targets. Grenioux et al, 2025.\n\n[2] Noble et al., Learned Reference-based Diffusion Sampling for Multi-Modal Distributions, ICLR, 2025.\n\n[3] Havens et al., Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching, ICML, 2025.\n\n[4] Chen et al., Sequential Controlled Langevin Diffusions, ICLR, 2025.\n\n[5] Towards Healing the Blindness of Score Matching, Zhang et al., 2022.\n\n[6] No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers. He \net al, 2025.\n\n[7] Improved Sampling via Learned Diffusions\n\n[8] Non-equilibrium Transport Sampler"}, {"review_id": "JJWldgHKy1", "reviewer": "Reviewer_TXNd", "summary": "This paper proposed a novel neural sampler based on annealed path and stochastic optimal control.\nUnlike previous works, which learn a single diffusion model or an interpolant-based process (e.g., based on escorted Jarzynski), this work is based on a combination of both. \nMore precisely, it learns two processes based on SOC. The first can be effectively trained with RAM and maps from a tractable prior to an intractable distribution. The second takes an interpolant process as the reference and maps the intractable distribution to the final target.", "questions": "1. Why is the ALDP exp only compared with DDS? What prior did you use for the ALDP experiment? Is the prior tuned to ensure coverage?\n\n\n2. Is it easy to learn to sample from the intractable prior? As shown in Fig 1, it does not make a big difference between the prior and the target. \n\n3. I still have difficulty intuitively understanding what makes NAAS better than other choices.\n\nFor example, intuitively, I suppose the 2nd transport (from an intractable prior to the final target) still has the teleportation issue, even though the prior is learnt, as the reference is defined by a geometric interpolant. But the vector field might be much easier compared to those defined by NETS/CMCD due to the similarity between the  intractable prior and the final target. \nOn the other hand, the first transport (from Gaussian to intractable prior) still has an uninformative reference. This might have an easier learning process compared to standard adjoint sampling. So can I conclude that NAAS performs better as it is an \"interpolant\" between AS and CMCD/NETS so that it can leverage both advantages?\n\nHowever, if this is the case, I am curious why this will not be a dilemma? More precisely, to make the first transport easier, it is better to have an intractable prior similar to the Gaussian. But this will lead to a more tricky learning procedure for the second transport. Is this a concern?\n\n\n4. Following the last question, if we concatenate two AS, or use $dX_t = -\\nabla U_t(X_t) dt + \\sqrt{2}dW_t, X_0 \\sim \\text{Gaussian}$ as the reference for the first process, will the results be worse than NAAS? \n \n\n5. I may not fully agree with the classification in the introduction. For example, both CMCD and PIS/DDS (or other SOC based sampler) are based on KL divergence between the current controlled process and the target process running from target marginal (can be viewed as optimal control). Why are PIS/DDS \"no reliance on IWS\" while CMCD is not?\n\n\n\n6. AS achieves high efficiency by decoupling gradient evaluations from direct calls to the target energy function. In contrast, NAAS requires explicit computation of the target in Eq 22. This bottleneck is also noted in recent literature, where many neural samplers are shown to underperform classical samplers like parallel tempering in terms of efficiency.  Is this also a limitation for NAAS? \n\n\n\n**This method is interesting and I generally enjoy reading this paper. I am very happy to raise my score to accept if the above question & concerns are addressed/discussed.**", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strength\n\n1. The method is novel and inspiring. Unlike previous works that focus on one family of reference processes, this approach combines two of them, which has the potential to address the limitation of both (for example, the Jarzynski-style reference), as it has information to guidance the sampler while the other one is easy to handle as it can be trained without simulation and has no teleportation issue. \n\n2. NAAS has promising performance compared to other baselines. It can also work on ALDP, a target that many of the current neural samplers failed without getting access to any data. \n\n3. The writing is very clear and easy to follow.\n\n## Weakness\n\n1. The motivation and classification seem to be inaccurate to me. For example, the authors motivate NAAS as a sampler that does not rely on importance sampling. However, it is hard to see why IS is a clear obstacle. It is easy to see the IS in iDEM suffers; however, I am not strongly convinced that the SMC in SCLD and PDDS has similar issue. Also, I do not understand why CMCD also relies on IS.\n\n2. I am not fully convinced why the proposed combination of two processes has advantages over other choices. Please see my questions for more details.", "was_revised": true, "final_justification": "The rebuttal addresses most of my concerns.  I am happy with the acceptance.  I believe the method itself is interesting and is well-grounded, but as there might be many modifications to the final paper, I keep my score as 4 due to a minor concern about how well these discussions (also discussions between authors and other reviewers) can be rearranged into the final version. Especially, I believe a stronger motivation would help to clarify “this is exactly what we need to do” rather than merely “stitching together what we can do”.", "rebuttal": "We appreciate the reviewer for spending time reading our manuscript carefully and providing valuable feedback. We hope that our responses are helpful in addressing the reviewer’s concerns. Due to the strict character limit, we have grouped and summarized related comments where appropriate.\n\n---\n**1. Clarification on motivation (W1)**\nWe agree that the current presentation of our motivation should be improved. Rather than positioning our method solely as an annealing-based sampler that avoids importance sampling, we will present it as one of the key characteristics of our approach. Instead, we would like to clarify that the core motivation and contribution of our work lies in developing a SOC-based sampler that incorporates annealed reference dynamics, as illustrated in the following flow:\n\n- Prior SOC-based approaches [1,2,3] have constructed unbiased diffusion samplers using stationary or uninformative base references, and do not fully leverage annealing references as employed in other importance sampling based methods. In this work, we propose a novel SOC-based framework that employs annealed reference dynamics as the base SDE.\n- Precisely, we formulate two SOC subproblems: (i) learning a prior distribution $\\mu$ from the initial distribution $\\delta_0$ with standard base reference, and (ii) transporting $\\mu$ to the target $\\nu$ using controlled annealed dynamics. We unify these into a single SOC formulation, which yields a controlled diffusion process that effectively transports $\\delta_0\\to\\mu\\to\\nu$.\n- Among available solvers for SOC, we adopt Adjoint Matching [3], a recently proposed method being scalable and sample-efficient. As discussed in Adjoint Sampling (AS) [4], this approach enables scalability and improved training stability.\n\n---\n**2. Clarification on importance sampling (W1)**\n\nAs established in prior works (e.g. [5]; [6] Section 4.1.4), importance sampling can suffer from high variance in high-dimensional settings, particularly when estimating expectations under complex distributions. In such cases, one must compute unnormalized weights $\\\\{w^{(i)}\\\\}^N_{i=1}$ corresponding to sampled trajectories $\\\\{X^{(i)}\\\\}^N_{i=1}$, where $$w^{(i)}\\propto\\frac{dp^\\star}{dp}(X^{(i)}),$$ where $p^\\star$ denotes the target path measure, while $p$ is the proposal. These weights are then used to form an empirical approximation of the target distribution,$$\\nu\\approx\\sum^N_{i=1}w^{(i)}\\delta_{X^{(i)}}.$$\n\nWhen the support or mass of the proposal and target distributions are misaligned, as is common in high-dimensional or complex energy landscapes. This leads to poor effective sample size and significant variance in the empirical approximation, which is further exacerbated when the number of samples $N$ is insufficient. For example, in PDDS [7], the learned distribution approximates a reweighted empirical measure and is thus exposed to the instability induced by high-variance weights. In contrast, our method learns a sampling distribution directly via path-space control, thereby avoiding the need to compute or rely on importance weights.\n\n---\n**3. Clarification on combination of two stochastic processes (W2, Q2, Q3, Q4)**\n\n**3.1 The role of two optimization stages (W2, Q2)**\n\nWe appreciate the reviewer for raising this important point. We agree that we need a more intuitive explanation of the role of two optimization problems. We incorporate the following intuitive explanation in our main manuscript in the last part of Section 3.\n\nThe two optimization problems in our framework, (i) learning the prior distribution $\\mu$ over the interval [−1,0], and (ii) optimizing the controlled annealed dynamics over [0,1], exhibit fundamentally different learning dynamics and objectives.\n\nThe first-stage optimization over [−1,0] learns a prior distribution $\\mu=e^{-U_0-V_0}$, where $U_0$ is a simple, prescribed energy function (typically corresponding to Gaussian distribution). This prior is designed to have broad and smooth support, in contrast to the often sharp or multimodal structure of the target distribution $\\nu$. The objective here is not to capture fine-scale structure, but to construct a low-complexity distribution that can be easily sampled using a simple base process, while still broadly covering the support of $\\nu$.\n\nIn contrast, the second-stage optimization over [0,1] focuses on refining the annealed dynamics to guide trajectories from $\\mu$ toward the target distribution $\\nu$. Since the annealed reference dynamics already bias the flow toward high-density regions of the target, the control at this stage primarily serves to correct residual mismatches. Because we only required to learn the residual mismatch, the second optimization problem is also easy to learn.\n\nBy decoupling these two objectives, the alternating scheme enables more stable and effective training. Each optimization stage solves a simpler and more focused subproblem: the first constructs a well-behaved initialization, and the second incrementally corrects the flow toward the target. This separation reduces optimization interference and improves both convergence and sample quality.\n\n**3.2 Explanation of two subproblems and the extreme case (Q3)**\n\nThank you for pointing out the important issue that “NAAS performs better as it is an interpolant between AS and Annealed-based Methods”. We largely agree with this intuition, and we elaborated it on the previous paragraphs. In short, the first optimization term is AS [3] with easier target $\\mu$, and the second optimization term is correction of the annealed dynamics. \n\nThis decomposition allows each stage to focus on a well-scoped subproblem and mitigates the difficulty of learning a single controlled dynamics. As the reviewer noted, this design introduces a natural tradeoff: If $\\mu$ is very close to a standard Gaussian, the learning burden is shifted to the second stage, resembling an annealed SOC approach. Conversely, if $\\mu$ is close to $\\nu$, the second stage becomes nearly trivial, reducing the method to something closer to AS [3]. But, anyhow, compared to both AS [3] and annealing-based methods, NAAS consistently gains from partitioning the task into structured subproblems. \n\nWe thank the reviewer for this valuable insight. We will include a discussion of investigating this trade-off as a promising avenue for future work in the conclusion of the revised manuscript.\n\n**3.3 Why not use AS twice or annealed dynamics for 1st segment (Q4)**\n\nConceptually, one might consider alternative designs, such as stacking two AS stages or employing annealed dynamics in the prior learning stage. However, our method is deliberately designed to ensure unbiased sampling, which imposes structural constraints on how the reference dynamics are chosen in each stage.\n\nIn the first optimization stage (over time interval [-1,0]), the goal is to generate a sample proportional to the prior distribution $\\mu$ via our controlled dynamics. To guarantee unbiasedness of our controlled dynamics, the reference dynamics in this stage must be memoryless, like PIS [1] and DDS [2] as described in L97-115. Annealed dynamics violate this memoryless condition, and therefore cannot serve as a valid reference in this stage. Consequently, we adopt standard Brownian motion as the reference dynamics during [-1,0].\n\nAdditionally, to empirically validate the effectiveness of using annealed dynamics as a reference in the second stage (over [0,1]), we compare our method against AS [3], which shares the same SOC formulation but uses a static base. As shown in the table below, our method achieves significantly better performance highlighting the practical benefit of the annealed reference dynamics. We attribute this improvement to the fact that the annealed dynamics guide samples toward high-density regions of the target distribution, thereby facilitating the learning of optimal control.\n\n|Sinkhorn| MW54 | GMM40 | MoS|\n|----|----|----|----|\n|AS|0.32$\\pm$0.06|18984.21$\\pm$62.12|2178.60$\\pm$54.82|\n|Ours|**0.10**$\\pm$0.0075 |**496.48**$\\pm$27.08|**394.55**$\\pm$29.35|\n\n---\n**4. Clarification on remaining questions (Q1,Q5,Q6)**\n\n**Q1. ALDP&DDS.** Thank you for the question. In addition to DDS [2], we also implemented and evaluated several SOC-based samplers, including PIS [1] and LV-PIS [11], on the ALDP experiment. However, we couldn’t obtain converging result for LV-PIS. We provide the full comparison results in the table below:\n\n|Method|$\\phi$|$\\psi$|$\\gamma_1$|$\\gamma_2$|$\\gamma_3$|E$W_2$|\n|----|----|----|----|----|----|----|\n|PIS|0.597 |0.952|0.453|0.498|7.038|5.918|\n|DDS|0.493 |0.154|4.095|0.111|0.150|5.467|\n|Ours|0.260|0.236|0.272|0.132|0.156|**1.076**|\n\nAs shown in the table, our method (NAAS) achieves the best performance across most of the evaluated metrics. Regarding the prior, we learn the prior distribution $\\mu$ through a dedicated SOC problem over the interval [-1,0].\n\n**Q5. CMCD and IS.** We apologize for misleading you and thank you for helping us identify this important mistake. We should’ve put CRAFT [8] as a reference instead of CMCD [10]. We modified the manuscript.\n\n**Q6 AS.** We agree that NAAS incurs additional computational cost compared to methods like AS [3], due to the need to compute the adjoint state via an ODE and evaluate gradients of the energy function during training (see Eq. (22)). This computational overhead is indeed a limitation. We added a discussion of this limitation to the conclusion of our manuscript.\n\n---\nReferences\n\n[1] Path Integral Sampler\n\n[2] Denoising Diffusion Samplers\n\n[3] Adjoint Sampling\n\n[4] Adjoint Matching\n\n[5] Sequential Monte Carlo Samplers\n\n[6] Free energy computations: A mathematical perspective\n\n[7] Particle Denoising Diffusion Sampler\n\n[8] Continual Repeated Annealed Flow Transport Monte Carlo\n\n[9] Sequential Controlled Langevin Diffusions\n\n[10] Controlled Monte Carlo Diffusions\n\n[11] Improved Sampling via Learned Diffusions"}]

Author rebuttal:

Meta-review: The paper introduces a stochastic optimal control-based diffusion method. It is based on two stages where the first stage learns a distribution easy to sample from a simple base, and the second refines it towards the target distribution. It shows improved performance against baselines on standard benchmarks, showing reduced variance with respect to importance sampling. The method is however based on previous work on the Adjoint Matching procedure, and is expensive to train since it requires to learn two networks and to solve a backward adjoint state.
The rebuttal introduced important clarifications on motivation and positioning of the work, as well as additional experiments, with all reviewers agreeing on the contributions made by this paper.
Decision: Accept (poster)