Saved in:
Bibliographic Details
Main Authors: Zhu, Huminhao, Wang, Fangyikang, Zhang, Chao, Zhao, Hanbin, Qian, Hui
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2401.14069
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910307826270208
author Zhu, Huminhao
Wang, Fangyikang
Zhang, Chao
Zhao, Hanbin
Qian, Hui
author_facet Zhu, Huminhao
Wang, Fangyikang
Zhang, Chao
Zhao, Hanbin
Qian, Hui
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2401_14069
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Neural Sinkhorn Gradient Flow
Zhu, Huminhao
Wang, Fangyikang
Zhang, Chao
Zhao, Hanbin
Qian, Hui
Machine Learning
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.
title Neural Sinkhorn Gradient Flow
topic Machine Learning
url https://arxiv.org/abs/2401.14069