QDSB: Quantized Diffusion Schr\"odinger Bridges

TL;DR

QDSB introduces a quantized approach to diffusion Schr"odinger bridges, enabling faster training with comparable sample quality by using anchor-quantized distributions and cell-wise sampling.

Contribution

It proposes a novel quantized method for Schr"odinger bridges that reduces computational cost while maintaining stability and quality.

Findings

QDSB matches existing baselines in sample quality.

QDSB requires substantially less training time.

The method's stability is theoretically supported by error bounds.

Abstract

Learning generative models in settings where the source and target distributions are only specified through unpaired samples is gaining in importance. Here, one frequently-used model are Schr\"odinger bridges (SB), which represent the most likely evolution between both endpoint distributions. To accelerate training, simulation-free SBs avoid the path simulation of the original SB models. However, learning simulation-free SBs requires paired data; a coupling of the source and target samples is obtained as the solution of the entropic optimal transport (OT) problem. As obtaining the optimal global coupling is infeasible in many practical cases, the entropic OT problem is iteratively solved on minibatches instead. Still, the repeated cost remains substantial and the locality can distort the global transport geometry. We propose quantized diffusion Schr\"odinger bridges (QDSB), which compute the endpoint coupling on anchor-quantized endpoint distributions and lift the resulting plan back to original data points through cell-wise sampling. We show that the regularized optimal coupling is stable w.r.t. anchor quantization, with an error controlled by the quality of the anchor approximation. In real-world experiments, QDSB matches the sample quality of existing baselines, requiring substantially less time. Code and data are available at github.com/mathefuchs/qdsb.