Title: Adjoint Schrödinger Bridge Sampler
Venue: NeurIPS 2025
Accepted: yes
Keywords: Boltzmann distribution

diffusion sampler

Schrödinger bridge
Abstract: Computational methods for learning to sample from the Boltzmann distribution—where the target distribution is known only up to an unnormalized energy function—have advanced significantly recently. Due to the lack of explicit target samples, however, prior diffusion-based methods, known as _diffusion samplers_, often require importance-weighted estimation or complicated learning processes. Both trade off scalability with extensive evaluations of the energy and model, thereby limiting their practical usage. In this work, we propose **Adjoint Schrödinger Bridge Sampler (ASBS)**, a new diffusion sampler that employs simple and scalable matching-based objectives yet without the need to estimate target samples during training. ASBS is grounded on a mathematical model—the Schrödinger Bridge—which enhances sampling efficiency via kinetic-optimal transportation. Through a new lens of stochastic optimal control theory, we demonstrate how SB-based diffusion samplers can be learned at scale via Adjoint Matching and prove convergence to the global solution. Notably, ASBS generalizes the recent Adjoint Sampling (Havens et al., 2025) to arbitrary source distributions by relaxing the so-called memoryless condition that largely restricts the design space. Through extensive experiments, we demonstrate the effectiveness of ASBS on sampling from classical energy functions, amortized conformer generation, and molecular Boltzmann distributions. Codes are available at https://github.com/facebookresearch/adjoint_samplers


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers how to modify a reference SDE in order to sample from an unnormalized density. Like prior literature, it proposes to correct the SDE with an additive drift that should be learnt (Eq 2). Overall, there are two ways of learning the drift in the literature:\n\n*(i) This drift is learnt by minimizing a trajectory-level cost (Eq 4) and a general terminal cost*\n\nComputationally, this problem can be efficiently solved using so-called adjoint sampling (Eq 14). \n \nTheoretically, the resulting SDE does *not* sample the target density unless some additional, restrictive assumptions are made (memoryless schedule and Dirac initialization). \n\n*(ii) The drift is learnt by minimizing a trajectory-level cost (Eq 3) under a terminal condition*\n\nComputationally, this classical problem called Schrodinger Bridge can be solved using an Iterative Proportional Fitting algorithm. \n\nTheoretically, the resulting SDE *does* sample the target density. \n\nThe authors identify a special terminal cost (Eq 9) so that (i) and (ii) are theoretically equivalent, and in particular the resulting SDE samples from the target density. Because the special terminal cost depends on an unknown quantity, the authors propose to learn it on-the-go (Eq 15). This way, the authors aim to use the computational efficiency of (i) to solve the problem in (ii).\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\nThe paper is very clearly written and proposes to alternate between adjoint sampling updates (Eq 14) and corrective updates (Eq 15) to learn a corrective drift to an SDE for sampling from an unnormalized density. The alternate optimization scheme (Eqs 14-15) is connected to the classical Iterative Proportional Fitting algorithm which is appreciated. The benchmarks with other diffusion samplers is are satisfactory and the experiments on synthetic and molecular problems are appreciated. \n\n## Weaknesses\n\nOne of the main difficulties in sampling from an unnormalized density is finding the modes. It would be nice to visually inspect the algorithm's ability to find the modes as in [1]: please see Q2 for a more detailed request. \n\nIt is also unclear how the ASBS algorithm's ability to find modes scales as a function of the dimensionality and of the number of modes. \n\n[1] L. Illing Midgley, V. Stimper, G. N. C. Simm, B. Scho ̈lkopf, and J. Miguel Hernandez-Lobato. Flow annealed importance sampling bootstrap. In Interna- tional Conference on Learning Representations (ICLR), 2023.\n\nQUESTIONS:\nQ1. The authors propose a nested optimization problem, where the outer loop updates follow Eq 14 and the inner loop updates follow Eq 15. How brittle is the interplay between these two updates? For example, in the inner loop which approximates Eq 15, how many iterations of the optimizer are required to get stable results? \n\nQ2. Could the authors illustrate their algorithm's ability to find the modes, by illustrating it on the synthetic 40-mode Gaussian mixture in 2D from Figure 2 of [1]? Please use an initial distribution that has the same density as one of the modes of the mixture (so as to not assume that the user knows where the other modes are). \n\nThe resulting plot could be reported in the Appendix and could compare only ASBS and a vanilla Langevin to get a visual understanding of how the authors' algorithm deals with the \"search problem\" (finding the target modes) as compared to a basic local exploration algorithm (Langevin). \n\nHere is an additional discussion of why this plot would be useful. While the results on molecular tasks are impressive, they do not provide a visual understanding of which modes are found by the algorithm. For example, is it those modes that are closest to the initial distribution that are found by ASBS? Similarly, Figure 2, while illustrative, is an example where the target distribution is symmetric and the initial distribution is initialized in the barycenter of the target modes, so even a Langevin process would be an efficient sampler.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper attempts to tackle the problem of sampling from a Boltzmann distribution.  They do so by building on the recent Adjoint Sampling work, noting that the memoryless condition -- an issue first recognized in the Adjoint Matching paper -- may be solved by employing a stochastic optimal control approach to the Schrödinger Bridge problem.  The authors finally demonstrate ASBS's efficacy on a number of sampling tasks.\n\nSTRENGTHS AND WEAKNESSES:\nI found the paper to be quite well written, that the theoretical results -- especially that the memoryless condition can be relaxed via the proposed SOC approach -- were significant, and that the experimental results showed ASBS's efficacy.  As such, I believe the paper merits an acceptance.\n\nI don't have many notes or questions regarding the theoretical section of the paper as I found the theoretical results to be satisfactory.  However, I have a few concerns regarding the paper's experimental validation.  Further, I feel that while ASBS performed well in the presented experimental settings, the experimental settings considered in the paper are not ones which significantly push the boundary on amortized samplers (e.g., ALDP in internal coordinates and the various synthetic energy functions have been solved for some time now).  I already feel the paper merits acceptance, but if the authors were to include further experiments showing that ASBS scales well enough to tackle problems like alanine tripeptide or larger molecules in, e.g., cartesian coordinates (a more challenging task due to how internal coordinates allow one to bypass problems with very sharp energy functions) I would be happy to raise my score to a strong accept.  Note that prior work (albeit it is a Boltzmann generator) in [1] have been able to successfully to these sorts of problems, so it would be interesting to see whether ASBS could scale to this.\n\nAs a side note, I believe that the RDKit baseline is not representative of the actual state-of-the-art performance of RDKit based conformer generation, see [2] which achieves superior performance to most deep learning based conformation sampling methods using some different settings of RDKit than the ETKDG baseline used in this paper. \n\nReferences: \n\n[1] Tan, Charlie B., et al. \"Scalable equilibrium sampling with sequential boltzmann generators.\" ICML (2025).\n\n[2] Zhou, Gengmo, et al. \"Do deep learning methods really perform better in molecular conformation generation?.\" arXiv (2023).\n\nQUESTIONS:\n1. How does ASBS perform compared to the improved RDKit baseline in the conformer generation task?\n\n2. Can ASBS scale up to any of the even larger scale experimental settings I described in the strengths/weaknesses section?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a novel method for sampling from unnormalized densities. The method extends the adjoint sampler to allow for using arbitrary source distributions while retaining the scalability of matching-based objectives. The authors formalize their method using stochastic optimal control (SOC). In particular, they propose a SOC formulation of the Schroedinger Bridge problem which is key to their new methodology. However, this formulation, or more specifically, the terminal cost of the SOC problem depends on the Schroedinger potential which is not available. As such, the authors propose an optimization scheme that alternates between learning the control via SOC and updating the Schroedinger potential. They show that this scheme is an instance of the IPF algortithm and therefore converges to the Schroedinger Bridge between source and target density. They validate the effectiveness of their method on several synthetic energy functions, Alanine dipeptide, and amortized conformer generation.\n\nSTRENGTHS AND WEAKNESSES:\n### **Strengths And Weaknesses**\n\nStrenghts\n\n- Elegant SOC formulation of the Schroedinger Bridge problem\n- Both SOC and corrector objectives are highly scalable\n- The method is very sample efficient, which is crucial for most real-world applications\n\nWeakesses\n\n- Performance on ALDP (Ramachandran plot) is not very convincing, especially since internal coordinates were used. There exist several works that obtain significantly better performance, see e.g. [1].\n- The paper is heavily based on adjoint sampling\n\n[1] Temperature-Annealed Boltzmann Generators\n\nQUESTIONS:\n- Why are LV-PIS and LV-DIS reliant on importance weights as stated in Table 1?\n- To the best of my knowledge, IPF only converges to the Schroedinger Bridge for $u_0 = 0.$ How valid is the pre-taining scheme introduced in Appendix C then?\n- Can the log normalizer be computed? It seems like the importance weights in (78) are intractable\n- Why did the authors use such small buffer sizes?\n- What symmetries did the authors use for ALDP? I could not find anything mentioned in the appendix. Also, why use internal coordinates and not Cartesian as done for the conformer generation? I reckon the harmonic prior would also help for the ALDP task\n- The authors write that they use Gaussian priors for all other targets. Did they also use the corrector matching algorithm in this case? This would not be necessary since there is also a SOC formulation for Gaussian priors using DDS. A comparison would be very interesting here. \n- In the Appendix it says ‘For SCLD, we do not use subtrajectory splitting’. So, this would be CMCD with LV loss then?\n- How much did the harmonic prior help for the synthetic energy functions?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The authors introduced the Adjoint Schrödinger Bridge Sampler (ASBS), which is a novel diffusion-based sampler for Boltzmann distributions that addresses general Schrödinger bridge sampling problems using only the target energy function.

STRENGTHS AND WEAKNESSES:
Strengths:

It is known that the SOC setup for fine-tuning is limited to the Dirac measure [1] [2] for the initial distribution.
The authors extended it to handle arbitrary distributions with memoryless source distributions by leveraging Schrodinger bridge between two marginal distributions. Schrodinger bridge (SB) should be the right tool here.

Extensive evaluations are conducted with promising performance.

[1] Theoretical guarantees for sampling and inference in generative models with latent diffusions

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

Weaknesses:

I understand that SB are complex to understand. Howevere, the writings are a bit sloppy and many important details are less evalborated, e.g.  
	1) Eq.(5) is quite important to understand why we should avoid memory setups and the authors at least should cite the adjoint matching paper. To be honest, even for the adjoint matching paper, some details are not detailed or derived properly as well, which makes it harder to understand. 
	2) eq(6), if p^base is a distribution that takes a joint distribution, then we should not reuse it a non-joint distribution; the intuition of setting g(x):=log p_1^base / nu(x) is not elaborated. 
	3) eq(14) eq(15) are not evalaborated clearly and it is crucial to understand the dynamics and why it works, please write it in more details.

Didn't check proof.

QUESTIONS:
See comments on Weaknesses.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-06