Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schr\"odinger Bridge Matching

TL;DR

This paper introduces ASBM, a novel generative modeling framework that improves sampling efficiency and stability by recovering optimal trajectories in high-dimensional data through a two-stage Schr"odinger Bridge approach.

Contribution

We propose a non-memoryless Schr"odinger Bridge-based method that enhances trajectory efficiency and scalability in high-dimensional generative modeling.

Findings

ASBM produces straighter, more efficient sampling paths.

It scales effectively to high-dimensional data with improved stability.

Experiments show better fidelity with fewer sampling steps.

Abstract

Diffusion models often yield highly curved trajectories and noisy score targets due to an uninformative, memoryless forward process that induces independent data-noise coupling. We propose Adjoint Schr\"odinger Bridge Matching (ASBM), a generative modeling framework that recovers optimal trajectories in high dimensions via two stages. First, we view the Schr\"odinger Bridge (SB) forward dynamic as a coupling construction problem and learn it through a data-to-energy sampling perspective that transports data to an energy-defined prior. Then, we learn the backward generative dynamic with a simple matching loss supervised by the induced optimal coupling. By operating in a non-memoryless regime, ASBM produces significantly straighter and more efficient sampling paths. Compared to prior works, ASBM scales to high-dimensional data with notably improved stability and efficiency. Extensive experiments on image generation show that ASBM improves fidelity with fewer sampling steps. We further showcase the effectiveness of our optimal trajectory via distillation to a one-step generator.