Adjoint Schrödinger Bridge Sampler

2025-06-27

AI Summary

Existing diffusion samplers for Boltzmann distributions with only unnormalized energy functions rely on target samples, importance weighting, or complex training—sacrificing scalability for estimation accuracy. This work proposes a scalable, target-sample-free diffusion sampling framework. We first integrate Schrödinger Bridge (SB) theory with generalized adjoint matching, relaxing the conventional memoryless assumption to accommodate arbitrary source distributions while ensuring global convergence. Grounded in stochastic optimal control and kinetic optimal transport, our approach yields an end-to-end differentiable SB-based diffusion model. Evaluated on classical energy-based models, conformal generation, and molecular conformation sampling, the method achieves significant gains in sampling efficiency and distribution fidelity. It bridges theoretical rigor—guaranteeing principled convergence—with practical scalability, enabling robust, gradient-based inference without requiring ground-truth samples or auxiliary reweighting schemes.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: User modeling for targeted and personalized online advertising

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.

Problem

• Sampling from Boltzmann distribution without target samples
• Improving diffusion sampler scalability and efficiency
• Generalizing adjoint sampling to arbitrary source distributions

Innovation

• Uses Schrödinger Bridge for efficient sampling
• Employs Adjoint Matching for scalable learning
• Generalizes Adjoint Sampling to arbitrary sources