Variational Schr\"odinger Diffusion Models

TL;DR

The paper introduces VSDM, a scalable variational inference-based diffusion model that improves transportation plan efficiency and training simplicity, with proven convergence and competitive results on CIFAR10 and time series data.

Contribution

We propose VSDM, a novel variational Schr"odinger diffusion model that linearizes forward scores, enabling simulation-free training and improved scalability in diffusion models.

Findings

VSDM efficiently generates anisotropic shapes with straighter trajectories.

VSDM achieves competitive unconditional generation on CIFAR10.

VSDM is tuning-friendly and does not require warm-up initializations.

Abstract

Schr\"odinger bridge (SB) has emerged as the go-to method for optimizing transportation plans in diffusion models. However, SB requires estimating the intractable forward score functions, inevitably resulting in the costly implicit training loss based on simulated trajectories. To improve the scalability while preserving efficient transportation plans, we leverage variational inference to linearize the forward score functions (variational scores) of SB and restore simulation-free properties in training backward scores. We propose the variational Schr\"odinger diffusion model (VSDM), where the forward process is a multivariate diffusion and the variational scores are adaptively optimized for efficient transport. Theoretically, we use stochastic approximation to prove the convergence of the variational scores and show the convergence of the adaptively generated samples based on the optimal variational scores. Empirically, we test the algorithm in simulated examples and observe that VSDM is efficient in generations of anisotropic shapes and yields straighter sample trajectories compared to the single-variate diffusion. We also verify the scalability of the algorithm in real-world data and achieve competitive unconditional generation performance in CIFAR10 and conditional generation in time series modeling. Notably, VSDM no longer depends on warm-up initializations and has become tuning-friendly in training large-scale experiments.