Generalized Schrödinger Bridge Matching

Guan-Horng Liu,Yaron Lipman,Maximilian Nickel,Brian Karrer,Evangelos A. Theodorou,Ricky T. Q. Chen

2024-04-18

Abstract:Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider a generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Generalized Schrödinger Bridge (GSB), appears prevalently in many scientific areas both within and without machine learning. We propose Generalized Schrödinger Bridge Matching (GSBM), a new matching algorithm inspired by recent advances, generalizing them beyond kinetic energy minimization and to account for task-specific state costs. We show that such a generalization can be cast as solving conditional stochastic optimal control, for which efficient variational approximations can be used, and further debiased with the aid of path integral theory. Compared to prior methods for solving GSB problems, our GSBM algorithm better preserves a feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scalability. We empirically validate our claims on an extensive suite of experimental setups, including crowd navigation, opinion depolarization, LiDAR manifolds, and image domain transfer. Our work brings new algorithmic opportunities for training diffusion models enhanced with task-specific optimality structures. Code available at <a class="link-external link-https" href="https://github.com/facebookresearch/generalized-schrodinger-bridge-matching" rel="external noopener nofollow">this https URL</a>

Machine Learning,Optimization and Control

What problem does this paper attempt to address?

The paper primarily aims to address the following issues:

1. **Generalizing the distribution matching problem in diffusion models**: The paper proposes a generalized distribution matching setup, namely the Generalized Schrödinger Bridge (GSB) problem, which allows for considering task-specific objective functions to describe marginal distributions when matching between two boundary distributions.

2. **Developing new matching algorithms**: The authors propose the Generalized Schrödinger Bridge Matching (GSBM), a new matching algorithm capable of handling the non-trivial state cost \(V_t\) in the aforementioned generalized distribution matching problem. This algorithm can better maintain the transport mapping from one distribution to another and ensure stability and scalability during the training process.

3. **Solving methods under the optimal control framework**: The paper views this problem as a Conditional Stochastic Optimal Control (CondSOC) problem and proposes a series of efficient approximate solving methods, including solving methods based on Gaussian path approximation and methods utilizing path integral theory to reduce bias.

4. **Empirical validation**: Through a series of experiments (including crowd navigation, opinion depolarization, LiDAR manifold modeling, and image domain transfer), the effectiveness and superiority of the proposed GSBM algorithm are demonstrated.

In summary, the main contribution of this paper is the proposal of a new matching algorithm—GSBM, to solve distribution matching problems with task-specific optimization structures, and its effectiveness in multiple application scenarios is validated through experiments.