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.

翻译：给定泛函的Wasserstein梯度流（WGF）在机器学习文献中得到了广泛应用。近年来，神经网络被用于逼近底层Wasserstein梯度流中的某些难处理部分，从而实现了高效的推理过程。本文提出了神经Sinkhorn梯度流（NSGF）模型，该模型参数化从给定源分布到目标分布的Wasserstein梯度流（相对于Sinkhorn散度）的时间变化速度场。我们在NSGF中采用速度场匹配训练方案，该方案仅需源分布和目标分布的样本即可计算经验速度场近似。理论分析表明，随着样本量趋于无穷大，经验近似的均场极限收敛于真实底层速度场。为提升高维任务中的模型效率，进一步设计了双阶段NSGF++模型：首先沿Sinkhorn流快速逼近图像流形（≤5次标准流评估），然后沿简单直线流对样本进行精化。基于合成数据集与真实世界基准数据集的数值实验支持了理论结果，并验证了所提方法的有效性。