We develop novel neural network-based implicit particle methods to compute high-dimensional Wasserstein-type gradient flows with linear and nonlinear mobility functions. The main idea is to use the Lagrangian formulation in the Jordan--Kinderlehrer--Otto (JKO) framework, where the velocity field is approximated using a neural network. We leverage the formulations from the neural ordinary differential equation (neural ODE) in the context of continuous normalizing flow for efficient density computation. Additionally, we make use of an explicit recurrence relation for computing derivatives, which greatly streamlines the backpropagation process. Our methodology demonstrates versatility in handling a wide range of gradient flows, accommodating various potential functions and nonlinear mobility scenarios. Extensive experiments demonstrate the efficacy of our approach, including an illustrative example from Bayesian inverse problems. This underscores that our scheme provides a viable alternative solver for the Kalman-Wasserstein gradient flow.

翻译：我们开发了基于神经网络的新型隐式粒子方法，用于计算具有线性和非线性迁移函数的高维Wasserstein型梯度流。主要思路是利用Jordan-Kinderlehrer-Otto (JKO)框架中的拉格朗日表述，采用神经网络近似速度场。我们借鉴连续归一化流场景中神经常微分方程(neural ODE)的公式，实现高效密度计算。此外，我们利用显式递推关系进行导数计算，大幅简化了反向传播过程。该方法在处理各类梯度流方面展现出通用性，可适配多种势函数和非线性迁移场景。大量实验验证了本方法的有效性，包括一个来自贝叶斯逆问题的示例。这突显了该方案为Kalman-Wasserstein梯度流提供了可行的替代求解器。