Neural Sinkhorn Gradient Flow
This work addresses the challenge of scalable distribution approximation in machine learning, offering an incremental improvement over existing neural network-based gradient flow methods.
The paper tackles the problem of approximating Wasserstein gradient flows for efficient inference by introducing Neural Sinkhorn Gradient Flow (NSGF), which parametrizes velocity fields using neural networks and achieves convergence with theoretical guarantees, demonstrating effectiveness on synthetic and real-world datasets with up to 5 neural function evaluations for quick manifold approach.
Wasserstein Gradient Flows (WGF) with respect to specific functionals have been widely used in the machine learning literature. Recently, neural networks have been adopted to approximate certain intractable parts of the underlying Wasserstein gradient flow and result in efficient inference procedures. In this paper, we introduce the Neural Sinkhorn Gradient Flow (NSGF) model, which parametrizes the time-varying velocity field of the Wasserstein gradient flow w.r.t. the Sinkhorn divergence to the target distribution starting a given source distribution. We utilize the velocity field matching training scheme in NSGF, which only requires samples from the source and target distribution to compute an empirical velocity field approximation. Our theoretical analyses show that as the sample size increases to infinity, the mean-field limit of the empirical approximation converges to the true underlying velocity field. To further enhance model efficiency on high-dimensional tasks, a two-phase NSGF++ model is devised, which first follows the Sinkhorn flow to approach the image manifold quickly ($\le 5$ NFEs) and then refines the samples along a simple straight flow. Numerical experiments with synthetic and real-world benchmark datasets support our theoretical results and demonstrate the effectiveness of the proposed methods.