Title: Non-equilibrium Annealed Adjoint Sampler
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion sampler

stochastic optimal control

annealed sampling

adjoint matching

Boltzmann density
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\n1. 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.**\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations and potential negative societal impact of their work", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n**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.\n\nLIMITATIONS:\nTo 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\nPlease 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.*\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nhis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-06"}