Non-equilibrium Annealed Adjoint Sampler

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.

1. Clarification on motivation (W1) We 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:

• 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.
• 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$.
• 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.

2. Clarification on importance sampling (W1)

As 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)}}.$

When 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.

3. Clarification on combination of two stochastic processes (W2, Q2, Q3, Q4)

3.1 The role of two optimization stages (W2, Q2)

We 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.

The 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.

The 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$.

In 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.

By 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.

3.2 Explanation of two subproblems and the extreme case (Q3)

Thank 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.

This 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.

We 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.

3.3 Why not use AS twice or annealed dynamics for 1st segment (Q4)

Conceptually, 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.

In 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].

Additionally, 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.

4. Clarification on remaining questions (Q1,Q5,Q6)

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:

As 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].

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.

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.

References

[1] Path Integral Sampler

[2] Denoising Diffusion Samplers

[3] Adjoint Sampling

[4] Adjoint Matching

[5] Sequential Monte Carlo Samplers

[6] Free energy computations: A mathematical perspective

[7] Particle Denoising Diffusion Sampler

[8] Continual Repeated Annealed Flow Transport Monte Carlo

[9] Sequential Controlled Langevin Diffusions

[10] Controlled Monte Carlo Diffusions

[11] Improved Sampling via Learned Diffusions

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.

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.

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:

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.

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.

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.

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.

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.

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.

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.

As 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$.

We 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.

To 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:

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)?

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.

W7. No indication of setting in realistic sampling frameworks is provided in the main paper, which hurts the applicability to real word problems

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.

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.

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$.

L1 Limitations Several drawbacks: (1) requires access to the Hessian of $U_t$ to apply Adjoint Matching, (2) requires two neural networks (see [6]).

A. We agree with the reviewer that we should discuss the limitation that pointed out. We will incorporate it in the conclusion section.

Once again, we appreciate the reviewer for careful consideration and important suggestions. We hope that our responses are helpful in addressing the reviewer’s concerns.

References

[1] Improving the evaluation of samplers on multi-modal targets. Grenioux et al, 2025.

[2] Noble et al., Learned Reference-based Diffusion Sampling for Multi-Modal Distributions, ICLR, 2025.

[3] Havens et al., Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching, ICML, 2025.

[4] Chen et al., Sequential Controlled Langevin Diffusions, ICLR, 2025.

[5] Towards Healing the Blindness of Score Matching, Zhang et al., 2022.

[6] No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers. He et al, 2025.

[7] Improved Sampling via Learned Diffusions

[8] Non-equilibrium Transport Sampler

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.

1. Clarification on motivation

Summary of reviewer’s comments: The 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.

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).

1.1 The roles of $u$ and $U_t$

The 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$.

However, 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.

1.2 Discussion on weight-based vs AM-based SOC solvers

In 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.

We 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.

We 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.

1.3 Why do we need to learn the prior distribution

The 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.

1.4 Clarification of our motivation & contribution

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:

• 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.
• (*) 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$.
• 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.

2. Clarification on reweighting

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.

Regarding 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.

3. Clarification on remaining questions

Q1. The meaning of “non-equilibrium”.

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$.

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.

A. We agree with the reviewer’s comment. We revised the table to use the more neutral label “Uses IWS during training”.

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.

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.

Q4. The notation used to index functions by time is not consistent.

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.

Q5. Typo / Minor comments on L110–111, L156, Corollary 3.2, Figure 1, Table 2, citation/reference.

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.

Once again, thank you for your careful review and thoughtful suggestions. We hope that the revisions we’ve made address your concerns.

[1] Non-equilibrium Transport Sampler

[2] Solving high-dim HJB PDEs using neural networks

[3] Adjoint Matching

[4] Adjoint Sampling

[5] Sequential Monte Carlo Samplers

[6] Free energy computations: A mathematical perspective

[7] Path Integral Sampler: a stochastic control approach for sampling

[8] Denoising Diffusion Samplers

[9] Sequential Controlled Langevin Diffusions

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.

1. Clarification on alternating vs joint optimization (W1, W2, Q2)

Summary of reviewer’s comments: The 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).

Author response: We 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. The 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.

Decoupling 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.

Finally, 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.

We 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.

2. Additional experiments on higher-dim problems (W3, W4, Q1)

Summary of reviewer’s comments: The reviewer noted a lack of empirical scalability analysis (W3), requested comparison on higher-dim problems with other baselines (Q1), and potential failure modes (W4).

Author response: We thank the reviewer for highlighting the importance of evaluating scalability and for suggesting a more detailed empirical analysis on high-dimensional settings.

While 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.

In 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:

Our 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.

Regarding 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.

3. Discussion on variance reduction (W4, Q3)

Author response: We 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.

Formally, 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.

For 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.

Additionally, 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.

4. Clarification on remaining questions

Q4: Missing SMC references

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.

Q5: Presentation

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 clarified the stability and scalability aspects of our method, Expanded the discussion on design choices (e.g., alternating vs. joint optimization), and Incorporated additional discussion on importance sampling and its variance as 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.

Q6: Generalization of NAAS beyond alanine dipeptide (ALDP) 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.

On 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.

References

[1] Improved Sampling via Learned Diffusions

[2] Non-equilibrium Transport Sampler

[3] Particle Denoising Diffusion Sampler

[4] Sequential Controlled Langevin Diffusions

[5] Sequential Monte Carlo Samplers

[6] Free Energy Computations: A Mathematical Perspective

[7] Continual Repeated Annealed Flow Transport Monte Carlo

[8] Path Integral Sampler

[9] Adjoint Sampling

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.