Numerical Analysis on Neural Network Projected Schemes for Approximating One Dimensional Wasserstein Gradient Flows

TL;DR

This paper develops a neural network-based numerical scheme for approximating one-dimensional Wasserstein gradient flows, providing theoretical guarantees and demonstrating effectiveness through various numerical examples.

Contribution

It introduces a novel neural projected scheme using ReLU networks and Wasserstein natural gradient with theoretical analysis and practical validation.

Findings

The scheme accurately approximates gradient flows in various models.

Theoretical guarantees ensure well-posedness and consistency.

Numerical examples confirm the method's effectiveness.

Abstract

We provide a numerical analysis and computation of neural network projected schemes for approximating one dimensional Wasserstein gradient flows. We approximate the Lagrangian mapping functions of gradient flows by the class of two-layer neural network functions with ReLU (rectified linear unit) activation functions. The numerical scheme is based on a projected gradient method, namely the Wasserstein natural gradient, where the projection is constructed from the $L^2$ mapping spaces onto the neural network parameterized mapping space. We establish theoretical guarantees for the performance of the neural projected dynamics. We derive a closed-form update for the scheme with well-posedness and explicit consistency guarantee for a particular choice of network structure. General truncation error analysis is also established on the basis of the projective nature of the dynamics. Numerical examples, including gradient drift Fokker-Planck equations, porous medium equations, and Keller-Segel models, verify the accuracy and effectiveness of the proposed neural projected algorithm.