Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling
Valentin De Bortoli, James Thornton, Jeremy Heng, Arnaud Doucet
NeurIPS 2021 (spotlight)
stat.ML, cs.LG, math.PR
2021-06-02
DSB runs neural IPF on the Schrödinger bridge. Iterate 1 is Song's SGM; later ones fix short-T mismatch. MNIST at 12 steps, CelebA 32x32 at 50 steps and 10 iterates, below SOTA.
Score-based generative models add noise until the data look like an easy prior, usually Gaussian, then learn the reverse denoising dynamics. Song et al. (2021) write the forward noising as an SDE whose reverse drift is the score (gradient of the log density) of the perturbed law, fit by a neural net.
The bottleneck is the time budget. The forward SDE has to run long enough for the terminal law to approach the Gaussian prior, and the step size has to be small enough for Euler-Maruyama to stay accurate. Together those constraints stretch the sampler. Early Song and Ermon (2019, 2020) setups used N = 100 steps.
A Schrödinger bridge asks a different question. Given a reference diffusion, data pdata and prior pprior, find the path-space process closest in KL to the reference whose time-0 and time-T marginals equal pdata and pprior. Sampling then hits both ends in finite time T, without waiting for the forward process to mix. The static form is entropy-regularized optimal transport; as the noise vanishes it recovers the quadratic Wasserstein plan.
The solver is Iterative Proportional Fitting (IPF). Start from the reference noising process. Odd iterates pin the terminal law to pprior; even iterates pin the initial law to pdata. On a finite discrete state space this is Cuturi's (2013) Sinkhorn algorithm.
DSB approximates each iterate's forward and reverse kernels by isotropic Gaussians whose means are two networks, F and B. Matching every historical score would cost too much memory, so training uses mean-matching: push samples with the current forward drift, then train the reverse net to recover where each step came from.
T can be short. Iteration 1 then fails because pN is not Gaussian; extra IPF rounds straighten the bridge.
Theorem 1 is the first quantitative TV bound for Song-style SGM. If the score is uniformly within M of the truth, the TV distance between generated samples and pdata splits into a discretization-plus-fitting term that grows like exp(D T), and a mixing term from the forward process: exp(-√α T) for OU (α>0), and T^{-1}+T^{-1/2} for Brownian motion. Mixing depends only on the forward Brownian or OU motion, not on dimension, so the usual high-dimensional Langevin mixing bottleneck does not appear. Larger α or T inflates the discretization constants, so steps must shrink.
Propositions 4 and 5 take IPF off compact support. Consecutive KL and TV gaps fall monotonically, n times the sum of the two marginal KLs goes to 0 (sharper than Léger's 2020 C/n bound), and the joints converge in TV. With extra density regularity the limit is the Schrödinger bridge. Gaussian references restore a geometric rate. DSB is the time discretization of this dynamic IPF on path space.
Gaussian sanity check: pdata = N(a,I), pprior = N(-a,I) with a = 0.1 times the all-ones vector, Brownian reference, closed-form static SB. Batch 128, N = 20, γ = 1/40. At d = 5 a 30k-parameter net matches mean, variance and covariance; at d = 50 the same architecture misses variance and covariance until the latent width is doubled to 240k parameters.
2D toys set T = 0.2 (N = 20, γk = 0.01). Iteration 1 is vanilla SGM: the endpoint is far from Gaussian and the reverse never returns to the data. In Figure 1 the two ends meet after 5 iterates; after 20 iterates synthesis is clearly better than iterate 1.
Images use a reduced U-Net from Nichol and Dhariwal (2021):
| Setting | N (horizon T) | DSB iterates | Result |
| 2D toys | 20 (0.2) | 1 vs 5/20 | Iterate 1 fails at short T; later iterates hit the data |
| MNIST | 12 | 8 | Visible gain over iterate 1; early SGM often used N = 100 |
| CelebA 32×32, d = 3072 | 50 (0.63) | 10 | Recognizable faces; claimed as the first practical high-d SB solver |
| EMNIST→MNIST | 30 (1.5) | 10 | Non-Gaussian prior, interpolation bridge |
FID falls with DSB iterates across several N. The paper only plots that curve in Figure 6 and never tables FID against DDPM or Score SDE. It states the implementation does not yet match then-SOTA.
Few-step sampling no longer requires the forward process to have mixed to a Gaussian. At short T, extra DSB rounds trade training iterates for sampling steps. MNIST at 12 steps and CelebA at 50, against early SGM's 100, is the right direction.
The prior need not be Gaussian. Swap in another dataset and the object is entropy-regularized transport between two data laws, a diffusion analogue of Sinkhorn on continuous state space. The paper flags multi-marginal SB, Wasserstein barycenters, entropy-regularized Gromov–Wasserstein, and domain adaptation. The solution is still a diffusion, so it can be rewritten as a neural ODE with the same marginals and exact likelihoods, as in Song et al.
This is a framework that attaches SGM to optimal transport. Image quality stops at generating samples and improving with extra iterates, not a new SOTA sampler. Speedups still need to be read next to noise schedules, better discretizations, and distillation.
T cannot go arbitrarily close to 0. When the forward endpoint sits far from pprior, there is almost no training signal near the prior's support, and the score remains hard to learn. The implementation also does not match SOTA.
Training is a sequence of nets. Later rounds can fine-tune earlier ones, but total compute is still well above a single score-matching run. If F and B are trained jointly, the equilibrium is some bridge, not necessarily the Schrödinger bridge.
Theorem 1 assumes a uniform score error bound M, which in practice one can at best hope for in high probability under the data distribution. The o(1/n) rate is for exact IPF; each DSB round adds network error and Euler error, and those three layers are never combined into a computable generation bound.
The image experiments omit a numeric FID table, and CelebA is only 32×32. Concurrent work by Vargas et al. (2021) estimates drifts with Gaussian processes; this paper uses neural nets and does not compare.