Variational Schr\"odinger Momentum Diffusion

TL;DR

This paper introduces Variational Schr"odinger Momentum Diffusion (VSMD), a scalable, simulation-free method for efficient generative diffusion that balances transport quality and computational cost.

Contribution

VSMD employs variational scores and a damping transform to improve scalability and stability, with proven convergence and superior empirical performance.

Findings

Efficiently generates anisotropic shapes with high transport quality.

Outperforms overdamped methods in various generative tasks.

Scales effectively to real-world data like time series and images.

Abstract

The momentum Schr\"odinger Bridge (mSB) has emerged as a leading method for accelerating generative diffusion processes and reducing transport costs. However, the lack of simulation-free properties inevitably results in high training costs and affects scalability. To obtain a trade-off between transport properties and scalability, we introduce variational Schr\"odinger momentum diffusion (VSMD), which employs linearized forward score functions (variational scores) to eliminate the dependence on simulated forward trajectories. Our approach leverages a multivariate diffusion process with adaptively transport-optimized variational scores. Additionally, we apply a critical-damping transform to stabilize training by removing the need for score estimations for both velocity and samples. Theoretically, we prove the convergence of samples generated with optimal variational scores and momentum diffusion. Empirical results demonstrate that VSMD efficiently generates anisotropic shapes while maintaining transport efficacy, outperforming overdamped alternatives, and avoiding complex denoising processes. Our approach also scales effectively to real-world data, achieving competitive results in time series and image generation.