Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adjoint Schrödinger Bridge Sampler

Guan-Horng Liu, Jaemoo Choi, Yongxin Chen, Benjamin Kurt Miller, Ricky T. Q. Chen

TL;DR

The paper tackles Boltzmann sampling when the target ν(x) ∝ e^{-E(x)} is only known up to an energy function. It reframes diffusion sampling as a Schrödinger Bridge problem and uses stochastic optimal control to derive a kinetic-optimal drift via Adjoint Matching, coupled with a Corrector Matching objective, enabling learning without target samples. By relaxing the memoryless condition, ASBS supports arbitrary priors and proves global convergence to the SB solution through an IPF-like alternating optimization. Empirically, ASBS outperforms multiple diffusion baselines on synthetic multi-particle energies, alanine dipeptide, and amortized conformer generation, demonstrating both effectiveness and scalability with domain priors. This approach offers a principled, computationally efficient route to Boltzmann sampling in chemistry and physics, with open-source code for reproducibility.

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. Code available at https://github.com/facebookresearch/adjoint_samplers

Adjoint Schrödinger Bridge Sampler

TL;DR

The paper tackles Boltzmann sampling when the target ν(x) ∝ e^{-E(x)} is only known up to an energy function. It reframes diffusion sampling as a Schrödinger Bridge problem and uses stochastic optimal control to derive a kinetic-optimal drift via Adjoint Matching, coupled with a Corrector Matching objective, enabling learning without target samples. By relaxing the memoryless condition, ASBS supports arbitrary priors and proves global convergence to the SB solution through an IPF-like alternating optimization. Empirically, ASBS outperforms multiple diffusion baselines on synthetic multi-particle energies, alanine dipeptide, and amortized conformer generation, demonstrating both effectiveness and scalability with domain priors. This approach offers a principled, computationally efficient route to Boltzmann sampling in chemistry and physics, with open-source code for reproducibility.

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. Code available at https://github.com/facebookresearch/adjoint_samplers

Paper Structure

This paper contains 25 sections, 8 theorems, 91 equations, 11 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Adjoint Schrödinger Bridge Sampler
  4. SOC Characteristics of the SB Problem
  5. Adjoint Sampling with General Source Distribution
  6. Alternating Optimization with Adjoint and Corrector Matching
  7. Theoretical Analysis
  8. Related Works
  9. Experiments
  10. Conclusion and Limitation
  11. Additional Preliminary
  12. Stochastic Optimal Control (SOC)
  13. Schrödinger Bridge (SB)
  14. Proofs
  15. Preliminary and Additional Theoretical Results
  16. ...and 10 more sections

Key Result

Theorem 3.1

The kinetic-optimal drift $u^\star_t$ in eq:u* solves an SOC problem

Figures (11)

  • Figure 1: Effect of the memoryless condition on learning SOC-based diffusion samplers. We consider Gaussian prior $\mu(x) := {\mathcal{N}}(x; 0,1)$ with $(f_t, \sigma_t)$ set to VP-SDE for the first plot and $(0,0.2)$ for the rest; see \ref{['appendix:soc']} for details. The memoryless condition injects significant noise (left) to correct the otherwise biased optimization (middle), whereas ASBS can successfully debias any non-memoryless processes (right).
  • Figure 2: Illustration of ASBS on a 2D example. By alternatively minimizing the Adjoint Matching (AM) objective \ref{['eq:am-obj']} and the Corrector Matching (CM) objective \ref{['eq:cm-obj']}, ASBS progressively learns a better corrector $h_\phi^{(k)}$ that debiases the SOC problem for the control $u_\theta^{(k)}$. Note that since the corrector is initialized with $h_\phi^{(0)}$$:= 0$, the first AM stage simply regresses $u_\theta^{(1)}$ to the energy gradient $\nabla E$.
  • Figure 3: Results on the synthetic energy functions for $n$-particle bodies with their corresponding dimensions $d$. Following chen2025sequentialAS, we report Sinkhorn for MW-5 and the Wasserstein-2 distances w.r.t samples, ${\mathcal{W}}_2$, and energies, $E(\cdot){\mathcal{W}}_2$, for the rest. All values are averaged over three random trials. Best results are highlighted.
  • Figure 4: The energy histograms of DW-4 and LJ-13 from \ref{['tab:synthetic']}. ASBS generates samples whose energy profiles closely match those of the ground-truth samples.
  • Figure 5: Complexity w.r.t. the number of function evaluation (NFE) on LJ-13 potential.
  • ...and 6 more figures

Theorems & Definitions (10)

  • Theorem 3.1: SOC characteristics of SB
  • Theorem 3.2: Global convergence of ASBS
  • Theorem 4.1: Adjoint Matching solves a forward half bridge
  • Theorem 4.2: Corrector Matching solves a backward half bridge
  • Lemma B.1: Itô lemma ito1951stochastic
  • Lemma B.2: Laplacian trick
  • proof
  • Theorem B.3: SB characteristics of SOC
  • proof
  • Corollary B.4: Reciprocal process of the SOC problem