Mean-field neural networks: learning mappings on Wasserstein space

TL;DR

This paper introduces neural network models designed to learn mappings between probability measures in Wasserstein space, supported by theoretical guarantees and demonstrated through numerical experiments on mean-field problems.

Contribution

It proposes two novel neural network architectures for mean-field functions with universal approximation guarantees, and applies them to solve time-dependent mean-field PDEs.

Findings

Neural networks accurately approximate mean-field functions.

Models generalize well across various test distributions.

Effective algorithms for solving mean-field PDEs are developed.

Abstract

We study the machine learning task for models with operators mapping between the Wasserstein space of probability measures and a space of functions, like e.g. in mean-field games/control problems. Two classes of neural networks, based on bin density and on cylindrical approximation, are proposed to learn these so-called mean-field functions, and are theoretically supported by universal approximation theorems. We perform several numerical experiments for training these two mean-field neural networks, and show their accuracy and efficiency in the generalization error with various test distributions. Finally, we present different algorithms relying on mean-field neural networks for solving time-dependent mean-field problems, and illustrate our results with numerical tests for the example of a semi-linear partial differential equation in the Wasserstein space of probability measures.