From Optimal Control to Optimal Transport via Stochastic Neural Networks in the Mean Field Setting

1. Introduction

In recent years, parametric Machine Learning (ML) applications, such as Neural Networks (NNs), have become central tools for various applications, ranging from computer vision and natural language processing to robotics and reinforcement learning. NNs consist of interconnected artificial neurons able to capture complex patterns while making accurate predictions. As their popularity has grown, also a fine mathematical description of the training procedure is increasingly required.

Along this latter research line, we propose a novel class of NNs, termed Mean Field Neural Networks (MFNNs), which are defined as the limiting object of a population of NNs when its number of components tends to infinity, as to derive a unified perspective for Mean Field Control (MFC) and Optimal Transport (OT) theory in an infinite dimension, allowing to get new insights about relationships between data in a concrete finite dimensional scenario.

We start the analysis by looking at the continuous idealization of a specific class of NNs, namely Residual NNs also named ResNets, whose training process is stated as a Mean Field Optimal Control Problem (MFOCP), and with a dynamic that evolves in terms of an ordinary differential equation (ODE). Moreover, the training problem of a ResNet is shown to be equivalent to a MFOCP of Bolza type, see [16] and [7] for further details.

Then we introduce a noise component into the dynamics of the ODE to derive a stochastic differential equation (SDE) that allows us to consider the inherent uncertainties connected to the variations in the real-world data thus allowing for the stochastic aspects of the learning process.

Finally, we consider the Supervised Learning problem analysis within the so-called Mean Field Optimal Transport (MFOT) formulation, recently introduced in [3]. We describe the MFC tools relevant to formalize the training process, hence we formulate the training problem as MFOT in an infinite-dimensional setting. It is worth mentioning that considering the collective interactions and distributions of the network’s parameters may facilitate the analysis of the network behaviour on a macroscopic level, hence improving the interpretability, scalability, and robustness of NNs models. Let us underline that the connection between mean field models and ML algorithms is also studied in [20], where the authors establish a mathematical relationship between the MFG framework and normalizing flows, a popular method for generative models composed of a sequence of invertible mappings. Similarly, in [9], the authors analyze Generative Adversarial Networks (GANs) from the perspectives of MFGs providing a theoretical connection between GANs, OT and MFG and numerical experiments.

The paper is organized as follows: in Section 2, we introduce the mathematical formalism of the supervised learning paradigm, while providing the description of the continuous idealization of a Residual NN stated as a MFOCP; in Section 3, we introduce a noisy component into the network dynamic, thus focusing on Stochastic NNs formalized as Stochastic Optimal Control problems; in Section 4 we review the MFG setting in a cooperative scenario defined in terms of MFC theory. Then, we consider recently developed Mean Field Optimal Transport methods that allow rephrasing MFC problems into OT ones. We also illustrate related approximation schemes and possible connection to an abstract class of NNs that respect the MFOT structure. We conclude by reviewing some methods to learn, i.e. approximate, mean field functions that depend on probability distribution obtained as the limiting object of empirical measures.

2. Residual Neural Networks as a Mean Field Optimal Control Problem

In this section, we present the workflow to treat a feed-forward NN, specifically a Residual NN, as a dynamical system based on the work in [23].

The main reference for this part is the well-known paper [16], where the authors introduce a continuous idealization of Deep Learning (DL) to study the Supervised Learning (SL) procedure which is stated as an optimal control problem by considering the associated population risk minimization problem.

2.1. The Supervised Learning paradigm

Following [15,24], the SL problem aims at estimating the function $\mathcal{F}:\mathcal{X}\to \mathcal{Y}$, commonly known as the Oracle. The space $\mathcal{X}$ can be identified with a subset of ${\mathbb{R}}^d$ related to input arrays (such as images, string texts or time series) while $\mathcal{Y}$ is the corresponding target set. Here for simplicity, we consider $\mathcal{X}$ and $\mathcal{Y}$ Euclidean spaces with different dimensions. Thus, training begins with a set of N input-target pairs ${\{x_0^i,y_T^i\}}_{i=1}^N$ where

• $x_0^i\in {\mathbb{R}}^d$ denotes the inputs of the NN;
• $x_T^i=\mathcal{F}\left(x_0^i\right)\in {\mathbb{R}}^d$ denotes the outputs of the NN;
• $y_T^i\in {\mathbb{R}}^l$ the corresponding targets.

We assume the same dimension of the Euclidean space for NN inputs and outputs allowing to explicitly write a dynamic in terms of a difference equation. Hence, for a ResNet, see [19] for more details, with T layers, the Feed-Forward propagation is given by

$$x_{t+1}=x_t+f(x_t,{\theta }_t)t=0,...,T-1$$

(1)

with $f:{\mathbb{R}}^d\to {\mathbb{R}}^d$ being a parameterized function and ${\theta }_t$ being the trainable parameters, e.g. bias, weights of layer t–th that belong to a measurable set $\mathcal{U}$ with values in a subset of the Euclidean space ${\mathbb{R}}^m$.

Remark 1. 

Following [30], we report an example of a domain for parameters of NN with ReLU activation functions. We define the following parameter domain

$$\mathsf{\Theta }=\left\{(a,w,b)\in \mathbb{R}\times {\mathbb{R}}^d\times \mathbb{R}:a^2<\left|\right|w\left|\right|+b^2\right\}$$

with activation functions $\varphi :\mathsf{\Theta }\to \mathbb{R}$ defined as

$$\varphi (\theta ;x)=a\sigma (w^{T}x+b),\theta =(a,w,b),\sigma \left(z\right)=z^{+}=max\{z,0\},$$

2.2. Empirical Risk Minimization

We aim at minimizing, over the set of measurable parameters $\mathsf{\Theta }$, a terminal loss function $\mathsf{\Phi }:\mathcal{X}\times \mathcal{Y}\to \mathbb{R}$ plus a regularization term L, to derive a Supervised Learning problem as an Empirical Risk Minimization (ERP) problem, namely

$$\underset{\theta \in \mathcal{U}}{min}\left[{\displaystyle \frac{1}{K}}\sum _{i=1}^N\mathsf{\Phi }(x_T^i,y^i)+\sum _{t=0}^{T-1}L\left({\theta }_t\right)\right]$$

(2)

over N training data samples indexed by i. We write $\theta =[{\theta }_0,...{\theta }_{T-1}]$ to identify the set of all parameters of the network.

If we consider no regularization of the parameters, i.e. $L=0$, and a quadratic loss function in terms of $\mathsf{\Phi }$, then Eq. (2) reads

$$J^{ERM}\left(\theta \right)=\underset{\theta \in \mathcal{U}}{min}\left[{\displaystyle \frac{1}{K}}\sum _{i=1}^N\left|\right|x_T^i-y^i{\left|\right|}^2\right]=\underset{\theta \in \mathcal{U}}{min}\left[{\displaystyle \frac{1}{K}}\sum _{i=1}^N\left|\right|f(x^i,\theta )-y^i{\left|\right|}^2\right]$$

(3)

being $x^i=[x_0,...,x_{T_1}]$ the discrete state process defined in Eq. (1).

Optimizing $J^{ERM}$ by computing its gradient is computationally expensive, especially if the number of data K is very large.

To handle the curse of dimensionality, it is usually common to initialize parameters from a ${\theta }^0$ from a probability distribution, to then optimize their choice inductively according to a Stochastic Gradient Descent scheme

$${\theta }^{k+1}={\theta }^k-{\eta }_t{\displaystyle \frac{1}{K}}\sum _{i=1}^K\left|\right|f(x^i,\theta )-y^i\left|\right|{\nabla }_{\theta }f(x^i,\theta )$$

(4)

with learning rate ${\eta }_t$ over K optimization steps.

For the sake of completeness, before going to the limit (we pass from a discrete set of training data to the corresponding distribution), we point out in the following remark that it is also possible to associate a measure corresponding to the empirical distribution of the parameters when the number of neurons goes to infinity.

Remark 2. 

A different approach, as illustrated, e.g., by Sirignano and Spilipupouls in [28], consists in associating to each layer the corresponding empirical measure and building a measure to describe the whole network, hence working with the empirical measure of controls, rather than states, as we see in Section 4. Following this perspective, the SGD equation (4) can be formalized as a minimization method over probability distribution. Moreover, the training of NN is based on the correspondence between the empirical measure of neurons ${\mu }_N$ and the function $f_N$ approximated by $NN$. Specifically, the training via gradient descent of a limit of an over-parameterized 1-hidden layer Neural Network with infinite width is equivalent to gradient flow in Wasserstein space [12,15,16,17]. Similarly, the convergence to solutions to SDEs is obtained as universal continuum objects in the small learning rate regime, see, e.g., [11].

From here on, we work with empirical distribution and measures associated with the training data.

2.3. Population Risk Minimization as Mean Field Optimal Control Problem

In what follows we move from the discrete setting to the corresponding continuous idealization by:

• going from layer index T to continuous parameter t;
• passing from a discrete set of inputs/output to distribution $\mu$ that represents the joint distribution in ${\mathcal{W}}_2({\mathbb{R}}^d\times {\mathbb{R}}^l)$ modelling the input-label distribution;
• passing from empirical risk minimization to population risk (i.e. minimization over expectation $\mathbb{E}$).

In particular, we pass to the limit in the number of data samples (number of input-target pairs) assuming also a continuous dynamic in place of layers discretization. The latter limit allows us to describe the dynamic of the state process x by the following Ordinary Differential Equation (ODE)

$${\dot{x}}_t=f(x_t,{\theta }_t),t\in [0,T],$$

(5)

in place of the finite difference equation (1). We identify the input-target pairs as sampled from a given distribution $\mu$ allowing us to write the SL problem as a Population Risk Minimization (PRM) problem.

In summary, we aim at approximating the Oracle function $\mathcal{F}$ using a provided set of training data sampled by a (known) distribution ${\mu }_0$ by optimizing weights ${\theta }_t$ to achieve maximal proximity between $x_T$ (output) and $y_T$ (target). Thus, we consider a probability space $(\mathsf{\Omega },\mathcal{F},\mathbb{P})$ and we assume inputs $x_0$ in ${\mathbb{R}}^d$ to be sampled from a distribution ${\mu }_0\in \mathcal{P}\left({\mathbb{R}}^d\right)$, with corresponding target $y_T$ in ${\mathbb{R}}^l$ sampled from a distribution $\nu \in \mathcal{P}\left({\mathbb{R}}^l\right)$, while the joint probability distribution $\mu$, that models the distribution of the input-target pairs, is defined by $\mu :=\mathcal{P}(x_0,y_T)$, belongs to the Wasserstein space ${\mathcal{W}}_2\left({\mathbb{R}}^{(d+l)}\right)$ and has ${\mu }_0$ and $\nu$ as its marginals. We recall that given a metric space $(X,d)$, the p-Wasserstein space ${\mathcal{W}}_p\left(X\right)$ is defined as the set of all Borel probability measures on X with finite p-moments.

The marginal distributions are obtained by projecting the joint probability distribution $\mu$ over the subspaces of inputs and output, respectively. We identify the first marginal, i.e. the projection over ${\mathbb{R}}^l$, with the distribution of inputs

$${\mu }_0={\int }_{R^l}\mu (x,y)dy,$$

while the distribution of targets reads

$$\nu ={\int }_{R^d}\mu (x,y)dx.$$

Moreover, we assume the controls ${\theta }_t$ depend on the whole distribution of input-target pairs capturing the mean-field aspect of the training data. We consider a measurable set of admissible controls, i.e. training weights, $\mathsf{\Theta }\subseteq {\mathbb{R}}^m$ and we state a Mean Field Optimal Control Problem (MFOCP) to solve the following PRM problem:

$$\underset{\theta \in L^{\infty }([0,T],\mathsf{\Theta })}{inf}{J}^{PRM}\left(\theta \right):={\mathbb{E}}_{\mu }\left[\mathsf{\Phi }\left(x_T,y_T\right)+{\int }_0^{T}L(x_t,{\theta }_t)dt\right]$$

(6)

$${\dot{x}}_t=f(x_t,{\theta }_t)0\le t\le T{x}_0\sim {\mu }_0,,(x_0,y_T)\sim \mu$$

We briefly report basic assumptions allowing us to have a solution for (6):

• $f:{\mathbb{R}}^d\times \mathsf{\Theta }\to {\mathbb{R}}^d$, $L:{\mathbb{R}}^d\times \mathsf{\Theta }\to \mathbb{R}$, $\mathsf{\Phi }:{\mathbb{R}}^d\times {\mathbb{R}}^l\to \mathbb{R}$ are bounded;
• f, L, $\mathsf{\Phi }$ are Lipschitz continuous with respect to x, with the Lipschitz constants of f and L being independent of parameters $\theta$, and
• $\mu$ has finite support in ${\mathcal{W}}_2\left({\mathbb{R}}^{(d+l)}\right)$

Problem (6) can be approached through two different methods: the first one is based on the Hamilton-Jacobi-Bellman (HJB) equation in the Wasserstein space, while the second one is based on a Mean Field Pontryagin Principle. We refer to [18] and [21] for viscosity solutions to the HJB equation in the Wasserstein space of probability measures, and to [8] for solving the constrained optimal control problems via Pontryagin Maximum Principle.

For the sake of completeness, let us also cite [6] where the authors introduce a BSDE technique to solve the related Stochastic Maximum Principle allowing us to consider the uncertainty associated with NN. The authors employ a Stochastic Differential Equation (SDE) in place of the ODE appearing in (6) to continuously approximate a Stochastic Neural Network (SNN). We deepen this approach in the next paragraph.

3. Stochastic Neural Network as a Stochastic Optimal Control Problem

In this paragraph, we generalize the previous setting considering a noisy dynamic, namely adding a stochastic integral to the deterministic setting described by the ODE in (6). The reference model corresponds to Stochastic NN whose discrete state process is described by the following equation

$$X_{n+1}=X_n+hF(X_n,{\theta }_n)+\sqrt{h}{\sigma }_n{\omega }_n,n=0,1,...,N-1$$

(7)

$\left\{{\omega }_n\right\}$ being a sequence of i.i.d. standard Gaussian random variables. We refer to [14] for a theoretical and computational analysis of SNN.

Eq. (7) can be generalized in a continuous setting. To this end, we consider a complete filtered probability space $\left(\mathsf{\Omega },\mathcal{F},{\mathbb{F}}^W,\mathbb{P}\right)$, and we introduce the following SDE

$$X_t=X_0+{\int }_0^{T}F\left(X_s,{\theta }_s\right)+{\int }_0^T{\sigma }_{s}d{W}_s,$$

(8)

with standard Brownian motion $W:={\left(W_t\right)}_{0\le t\le T}$ and diffusion term $\sigma$. Analogously as for ResNets, the index $T>0$ represents a continuous parameter modelling the width of the layer, being $X_T$ the output of the network.

Here, we report the theory developed in [2] to study Eq. (8) in the framework of SOC problem by introducing the control process $u=[\theta ,\sigma ]$. Thus we consider also the diffusion $\sigma$ as a trainable parameter of the model. We start by translating the SDE (8) into the following controlled process, written in differential form

$$d{X}_t=f(X_t,u_t)dt+g\left(u_t\right)d{W}_t,0\le t\le T,\left(3\right)$$

(9)

where $f(X_t,u_t)=F(X_t,{\theta }_t)$ and $g\left(u_t\right)={\sigma }_t$. As in classical control theory applied to ML, the aim is to select the control u that minimizes the discrepancy between the SNN output and the data. Accordingly, we define the cost function for our stochastic optimal control problem as

$$J\left(u\right):=\mathbb{E}\left[\mathsf{\Phi }\left(X_T,\mathsf{\Lambda }\right)\right],$$

(10)

$\mathsf{\Lambda }$ being a random variable that corresponds to the target of a given input, i.e. $X_0$. Then the optimal control $u^{★}$ is the one that solves

$$J\left(u^{★}\right)=\underset{u\in \mathcal{U}[O,T]}{inf}J\left(u\right)$$

above the class of measurable control $\mathcal{U}$. Always in [2], the authors deepen the Stochastic Maximum Principle approach to solve the latter problem. Firstly, the functional J is differentiated with respect to the control with a derivative in Gateaux sense over $[0,T]$

$$J_u^{\prime }(t,u_t)=\mathbb{E}\left[f_u^i{(X_t,u_t)}^T{Y}_t+g_u^{\prime }{\left(u_t\right)}^T{Z}_t\right].$$

(11)

Then, by the martingale representation of $Y_t$, the following Backward SDE is introduced

$$d{Y}_t=f_x^i{(X_t,u_t^{★})}^T{Y}_t+Z_{t}d{W}_t,Y_T={\mathsf{\Phi }}_x^{\prime }(X_T,\mathsf{\Lambda })$$

(12)

to model the back-propagation of the forward state process Eq. defined in (9) associated with the optimal control $u^{★}$.

Finally, the problem is solved by the gradient descent method with step-size ${\eta }_k$

$$u_t^{k+1}=u_t^k-{\eta }_k{J}_u^{\prime }(t,u_t^k),k=0,1,2,...,0\le t\le T1,.$$

(13)

The authors also provide a numerical scheme whose main benefit is to derive an estimate of the uncertainty connected to the output of this stochastic class of NNs.

We remark that here it is not possible to write the chain rule for Eq. (13) as previously done for Eq. (4) since the presence of the stochastic integral term that, differently from classical ML theory, makes the back-propagation itself a stochastic process, see Eq. (12). However, modern programming libraries (e.g., TensorFlow or PyTorch) perform the computation (13) automatically reducing the computational cost, hence allowing to go towards a mean field formulation (in terms of multiple interacting agents) of previous problems.

4. Mean Field Neural Network as a Mean Field Optimal Transport

As seen in Section 3, SOC deals with finding the optimal control policy for a dynamic system in the presence of uncertainty. Conversely, OT theory focuses on finding the optimal way to transport from one distribution to another.

In what follows, we recall some basic definitions related to the OT problem. Given two marginal distribution $\mu \in \mathcal{P}\left({\mathbb{R}}^d\right)$ and $\nu \in \mathcal{P}\left({\mathbb{R}}^d\right)$, the classical OT problem in the Kantorovich formulation reads

$$\underset{\pi \in \mathsf{\Pi }(\mu ,\nu )}{inf}\int c(x.y)\pi (dx,dy)$$

(14)

where c is a cost function and $\mathsf{\Pi }(\mu ,\nu )$ corresponds to the set of couplings between $\mu$ and $\nu$.

We focus on the setting where $\mu$ and $\nu$ are distributions computed on ${\mathbb{R}}^d$, i.e. $\mu \sim \left(X_1,...,X_d\right)$ and $\nu \sim \left(Y_1,...,Y_d\right)$. The Monge formulation reads

$$\underset{T:T\#\mu =\nu }{inf}\int c(x,T\left(x\right))\mu \left(dx\right)$$

(15)

where the infimum is computed over all measurable maps $T:{\mathbb{R}}^d\to {\mathbb{R}}^d$ with the pushforward constraint $T\#\mu =\nu$.

The possibility to link a SOC problem (and the related mathematical formulation of a specific learning procedure) to the corresponding OT formulation, allows for a deeper understanding of the underlying learning dynamics as well as the use of specific optimal transport techniques to analyze and solve the combined problem.

In this section, we focus on the connection between SOC and OT. This equivalence is particularly evident in the context of Mean Field Games (MFG), a class of stochastic control problems where a large number of agents interact and influence each other, and in particular the variational formulation of MFG which is directly linked to the dynamic formulation of OT by Benamou and Brenier, see, e.g., [4] for further details.

We mention that there are also specific scenarios where the dynamics of the stochastic control problem can be interpreted as a mass transportation problem, provided that certain assumptions on functionals and cost are guaranteed. For example in [25] and similarly in [29], the authors focus on an OT problem, whose cost depends on the drift and the diffusion coefficients of continuous semimartingale, and the minimization is run among all continuous semimartingales with given initial and terminal distributions.

4.1. Mean Field Game

Equations (5) and (8) look at the state process $X_t$ as a random variable, while the functional (10) is computed over the realization of this random variable without focusing on the interaction during the evolution. One further step relies on extending the previous equation to a McKean-Vlasov setting where the dynamic of a random variable X depends on the other N random variables by the mean of the distribution.

Adding the dependence on a mean-field term into the drift allows us to model the shared connections between the neurons.

We introduce the following McKean-Vlasov SDE for N particles/agents

$$X_t^i=X_0^i+{\int }_0^{T}b\left(X_s^i,m_{X_s}^N,{\theta }_s\right)+{\int }_0^T\sigma d{W}_s,i=1,...,N$$

(16)

with $X_0^i$ being the initial states. We assume a measurable drift $b:[0,T]\times {\mathcal{W}}_2\left({\mathbb{R}}^d\right)\times {\mathbb{R}}^d\to \mathbb{R}$, a constant diffusion $\sigma$ and we define the empirical distribution $m_{X_s}^N$ as

$$m_{X_s}^N={\displaystyle \frac{1}{N}}\sum _{j=1}^N{\delta }_{X_s^i}.$$

(17)

Going at the limit $N\to \infty$, the empirical measure $m^N$ tends to the probability measure m belonging to the Wasserstein space $W_2\left({\mathbb{R}}^d\right)$, i.e. the space of probability measures on ${\mathbb{R}}^d$ with a finite second–order moment.

More precisely, we introduce the following setting, that we need to define the solution of a MFG.

• a finite time horizon $T>0$;
• $\mathcal{Q}\subseteq {\mathbb{R}}^d$ is the state space;
• ${\mathcal{W}}_2\left(\mathcal{Q}\right)$ the space of probability measure over $\mathcal{Q}$;
• $(x,m,\alpha )\in \mathcal{Q}\times {\mathcal{W}}_2\left(\mathcal{Q}\right)\times {\mathbb{R}}^k$ describes the agent state, the mean-field term and the agent control;
• $f:\mathcal{Q}\times {\mathcal{W}}_2\left(\mathcal{Q}\right)\times {\mathbb{R}}^k\to \mathbb{R}$, $(x,m,\alpha )\mapsto f(x,m,\alpha )$ and $g:\mathcal{Q}\times {\mathcal{W}}_2\left(\mathcal{Q}\right)\to \mathbb{R}$, $(x,m)\mapsto g(x,m)$ provides the running and, resp., the terminal cost;
• $b:\mathcal{Q}\times {\mathcal{W}}_2\left(\mathcal{Q}\right)\times {\mathbb{R}}^k\to {\mathbb{R}}^d$ represents the drift function ;
• $\sigma >0$ is the volatility of the state.

Definition 1 

(MFG equilibrium). We consider a MFG problem with a given initial distribution $m_0\in {\mathcal{W}}_2\left(\mathcal{Q}\right)$. A Nash equilibrium is a flow of probability measures $\widehat{m}={\left(\widehat{m}(t,·)\right)}_{0\le t\le T}$ in ${\mathcal{W}}_2\left(\mathcal{Q}\right)$ plus a feedback control $\widehat{\alpha }:[0,T]\times \mathcal{Q}\to {\mathbb{R}}^k$ satisfying the following two conditions:

1.

$\widehat{\alpha }$ minimizes $J_m^{MFG}$over α:

$$\mathbb{E}\left[{\int }_0^{T}f(X_t^{m,\alpha },m(t,·),\alpha (t,X_t^{m,\alpha }))dt+g(X_T^{m,\alpha },m(T,·))\right]$$

where $\left(X_t^{m,\alpha }\right)$ solves the SDE

$$d{X}_t^{m,\alpha }=b(X_t^{m,\alpha },m(t,·),\alpha (t,X_t^{m,\alpha }))dt+\sigma d{W}_t$$

W being a d-dimensional Brownian motion and $X_0^{m,\alpha }$ has distribution $m_0$;

2.

for all $t\in [0,T]$, $\widehat{m}$ is the probability distribution of $X_t^{\widehat{m},\widehat{\alpha }}$.

4.2. Mean Field Control

Differently from MFG where players are modelled as competitors, MFC is a framework that considers a large population of agents aiming to cooperate and optimize their individual objectives while taking into account the collective behaviour of the entire population. In MFC, each agent seeks to minimize its own cost function by considering the impact of the mean field representing the average behaviour of all agents, hence capturing cooperation spirit among agents. Accordingly, the solution of a MFC is defined in the following way:

Definition 2 

(MFC optimum). Given $m_0\in {\mathcal{W}}_2\left(\mathcal{Q}\right)$, a feedback control ${\alpha }^{*}:[0,T]\times \mathcal{Q}\to {\mathbb{R}}^k$ is an optimal control for the MFC problem if it minimizes over α$J^{MFC}$ defined by

$$\mathbb{E}\left[{\int }_0^{T}f(X_t^{\alpha },m^{\alpha }(t,·),\alpha (t,X_t^{\alpha }))dt+g(X_T^{\alpha },m^{\alpha }(T,·))\right]$$

(18)

where $m^{\alpha }(t,·)$ is the probability distribution of the law of $X_t^{\alpha }$, under the constraint that the process ${\left(X_t^{\alpha }\right)}_{t\in [0,T]}$ solves the following SDE of McKean-Vlasov type

$$d{X}_t^{\alpha }=b(X_t^{\alpha },m^{\alpha }(t,·),\alpha (t,X_t^{\alpha }))dt+\sigma d{W}_t$$

(19)

with $X_0^{\alpha }$ having distribution $m_0$.

We refer to [10] for an extensive treatment of McKean-Vlasov control problems (18).

By considering the joint optimization problem of the entire population, MFC enables the analysis of large-scale systems with cooperative agents and provides insights into the emergence of collective behaviour. One possibility relies on stating the dynamic in Eq. (6) in terms of probability measures. For example, we can consider a continuity equation such as the Fokker-Planck equation to consider the evolution of the density function. Along this setting, we cite the measure-theoretical approach for NeurODE developed in [7] where the authors introduce a forward continuity equation in the space of measures with a constrained dynamic in the form of an ODE. Conversely, within the cooperative setting, we can also rely on a novel approach, named Mean Field Optimal Transport, introduced in [3], that we explore in the next paragraph.

4.3. Mean Field Optimal Transport

The Mean Field Optimal Transport deals with a framework where all the agents cooperate (such as in MFC) in order to minimize a total cost without terminal cost but with an additional constraint since also the final distribution is prescribed. We notice that the setting with fixed initial and terminal distributions resembles the one introduced in the Population Risk minimization problem described in Section 2. We follow the numerical scheme introduced in section 3.1 in [3] to approximate feedback controls, namely we introduce the following model.

Definition 3 

(Mean Field Optimal Transport). Let ${\mathbb{R}}^d$, describe the state space and denote by ${\mathcal{W}}_2\left({\mathbb{R}}^d\right)$ the set of square-integrable probability measures on ${\mathbb{R}}^d$. Let $f:{\mathbb{R}}^d\times {\mathcal{W}}_2\left({\mathbb{R}}^d\right)\times {\mathbb{R}}^k\to \mathbb{R}$ be the running cost, $g:{\mathbb{R}}^d\times {\mathcal{W}}_2\left({\mathbb{R}}^d\right)\to \mathbb{R}$ be the terminal cost, $b:{\mathbb{R}}^d\times {\mathcal{W}}_2\left({\mathbb{R}}^d\right)\times {\mathbb{R}}^k\to {\mathbb{R}}^d$ the drift function and $\sigma \in \mathbb{R}$ the non-negative diffusion. Given two distributions ${\rho }_0$ and ${\rho }_T\in {\mathcal{W}}_2\left({\mathbb{R}}^d\right)$ the aim of MFOT is to compute the optimal feedback control $v:[0,T]\times {\mathbb{R}}^d\to {\mathbb{R}}^m$ minimizing

$$J^{MFOT}:v\mapsto \mathbb{E}\left[{\int }_0^{T}f(X_t^v,{\mu }^v\left(t\right),v(t,X_t^v))dt\right]$$

(20)

being ${\mu }^v\left(t\right)$ the distribution of process $X_t^v$ whose dynamics is given by

$$\left\{\begin{array}{c}{X}_0^v\sim {\rho }_0{X}_T^v\sim {\rho }_T\hfill \\ d{X}_t^v=b(X_t^v,{\mu }^v\left(t\right),v(t,X_t^v))dt+\sigma d{W}_t,t\in [0,T]\hfill \end{array}\right.$$

(21)

with ${\rho }_0$ and ${\rho }_T$ are the prescribed initial and terminal distributions.

This type of problem incorporates mean field interactions in the drift and the running cost. Furthermore, it encompasses classical OT as a special case by considering $b(x,\mu ,a)=a$, $f(x,\mu ,a)=\frac{1}{2}{a}^{T}a$ and $\sigma =0$.

The integration of MFC and OT allows to both tackle the weight optimization problem in NN, and to model the flow of information or mass between layers of neurons, while the optimal weights may be computed as the minimizers of the functional with respect to controls v

$$v^{★}=\underset{v\in \mathcal{U}}{min}{J}^{MFOT}\left(v\right)$$

(22)

along all the trajectories $X^v$,being $\mathcal{U}$ the set of admissible controls.

Thus, we look at the MFNN as a collection of identical, interchangeable, indistinguishable NNs where the dynamic of the representative agents is a generalization of a SNN (7) allowing a dependence on the term ${\mu }^v\left(t\right)$ modelling the mean field interactions. By considering the MFNN dynamic as a population of interconnected NNs, we can employ mean-field control to analyze the collective behavior and interactions of networks, accounting for their impact on the overall network performance.

To summarize, we are looking at this novel class of NN, i.e. MFNN, as the asymptotically configuration of NNs in a cooperative setting.

We remark that the representative agent does not know the mean field interaction terms since it depends on the whole population but an approximated version can be recursively learned. For example in [3] the authors present different numerical scheme to solve MFOT:

1.

Optimal control via direct approximation of controls v;

2.

Deep Galerkin Method for solving a forward-backward systems of PDEs;

3.

Augmented Lagrangian Method with Deep Learning exploiting the variational formulation of MFOT and the primal/dual approach.

We briefly review the direct method (1) to approximate controls of feedback type by an optimal control formulation. The controls are assumed of feedback form and can be approximated by

$$g(x,\mu )=G\left({\mathcal{W}}_2(\mu ,{\rho }_T)\right),\mu \in {\mathcal{W}}_2\left({\mathbb{R}}^d\right)$$

(23)

where $G:{\mathbb{R}}_{+}\to {\mathbb{R}}_{+}$ is an increasing function. The idea is to use the function in Eq. (23) as a penalty for being far from the target distribution ${\rho }_T$ as the terminal cost to embed the problem into classical MFG/MFC literature. Intuitively, Eq. (23) corresponds to the infinite dimensional analogous of the loss function of the leveraged NN algorithm, being $\mu$ the final distribution that has to be close as possible in the sense of Wasserstein metric to the target distribution ${\rho }_T$.

In view of obtaining a numerically tractable version of the SDE (21), one may consider a classical discretization Euler-Maruyama scheme, also requiring the set of controls v to be restricted to the ones approximated by a NNs $v_{\theta }$ with parameters $\theta$. Moreover, approximating the mean field term m by its finite dimensional counterpart, see Eq. (17), allows to develop a stable numerical algorithm, see Section 3.1 in [3] for further details, particularly w.r.t. the linked numerical implementation.

4.4. Other approaches for learning Mean Field function

For the sake of completeness we also mention two different methods to deal with the approximation of mean field function that can be used in parallel with MFOT:

• the first data-driven approach, presented in [1], has been considered to solve a stochastic optimal control problem, where the unknown model parameters were estimated in real-time using a direct filter method. This method involves transitioning from the stochastic maximum principle to approximate the conditional probability density functions of the parameters given an observation, which is a set of random samples;
• in [26], the authors report a map that by operating over an appropriate classes of neural networks, specifically the Bin density-based approximation and Cylindrical approximation, is able to reconstruct a mapping between the Wasserstein space of probability measures and an infinite dimensional function space on a similar setting to MFG.

5. Conclusions and further directions

In the present article, we provided an all-around overlook of methods at the intersection of parametric ML, MFC and OT. By assuming a dynamical system viewpoint we considered the deterministic, ODE–based, setting of the supervised learning problem, to then incorporate noisy components, allowing for stochastic NNs definition, hence introducing the MFOT approach. The latter, derived as the limit in the number of training data, recasts the classical learning process in a mean field optimal transport one. As a result, we gained a unified perspective for the parameters optimization process characterizing ML models with a specified learning dynamic, within the framework of OT and MFC, which may concede to efficiently handle high dimensional data sets.

We empathise that the major limitation of MFOT (20) concerns the fact that many of its convergence results, such as those related to corresponding forward–backward systems, still need to be verified. Nevertheless, it represents an indubitably fertile and stimulating research ground that should be enhanced since it permits the derivation of techniques that may significantly improve the robustness of algorithms, particularly when dealing with huge sets of training data potentially perturbed by random noise components.