Neural Networks and Adapted Optimal Transport

1. Introduction

The present paper aims at presenting a baseline for the problem of including causality constraint into Machine Learning (ML) tasks as to increase algorithms robustness, scalability and accuracy, particularly when dealing with structured time series-based applications. The analysis starts from a preliminary study of a Supervised Learning problem in terms of a Mean Field Optimal Control Problem (MFOCP), hence providing an overview about the most relevant theoretical aspects that can be useful for a generalization of MFOCP to include the temporal structure of the training data. Accordingly, we provide a complete review about Adapted Optimal Transport (AOT), that is a version of Optimal Transport (OT) enforcing an adapted (with respect to a filtration) constraint. Furthermore, we summarizes state-of-the-art works at the crossroad of probability theory, ML and OT methods. We deepen the introduction of probability measures into ML tasks specifically for Neural Networks (NNs) applications in the context for explainable AI. On the other hand, the use of probability measures can also help to analyze the behavior of specific algorithm, while NNs can be used to approximate functions depending on unknown distributions. Trying to clarify the interplay between probability measures and Neural Networks, we recall possible cross-related cases:

• Measure approach for NN training: Probability measures can be used to analyze the behavior of neural networks, both during training and inference. For example, the loss function used to train a neural network can be seen as a divergence between the model distribution and the target distribution, which can be expressed in terms of probability measures, [17]. Similarly, the output of a neural network can be interpreted as a probability distribution over the output space, which can be analyzed using probability theory. This can be useful for understanding the uncertainty of the predictions made by the neural network, such as the Stochastic Deep Networks (SDN) introduced in [25]. The design of a SDN involves a deep architecture that can handle probability distribution inputs and outputs and uses the classical Wasserstein distance as the loss function. Moreover, layers correspond to a sequence of elementary blocks that maps random vectors to random vectors to capture multiple interactions between samples from the distributions.
• Probability measures as a tool to regularize neural networks: Probability measures can be used as a tool to regularize neural networks, by adding a regularization term to the loss function that encourages the output of the neural network to match a prior distribution by the so called entropic measure, such as in Eq. (26). Within this case, we mention the entropic version of Sinkhorn algorithm where an entropic regularization is added, developed, e.g., in [40]. This algorithm can be useful for scaling OT problem in high dimensional setting or for generative modeling tasks, where the goal is to learn a distribution that can generate new data samples.
• Neural networks to approximate empirical measures: Neural networks can be used to approximate probability measures, such as empirical distributions over a state space for mean field function. We specifically consider this feature in Section 3. Analogously, one can consider NN architecture for learning population dynamics. For example, in [33], the authors focus on a recurrent architecture to model the diffusion of a population by means of a NN by injecting random noise.
• Neural Network to learn unknown probability measure for unsupervised learning tasks such as density estimation, where the goal is to learn the underlying distribution of the data. For example, in [21], the authors try to build an algorithm to directly learn a probability measure via an optimal transport metric, also providing some convergence results for the learning algorithm.

The paper is organized in the following paragraphs. In Section 2, we review some basic works dealing with the mathematical formalism and definitions needed to precisely address supervised learning problems within the context of NNs. In Section 3, we review some developed methods to learn, i.e. approximate, mean field functions that depend on probability distribution obtained as limit object of empirical measures. In Section 4, we deal with the problem of optimizing over the set of probability measures with finite 2-moment focusing on JKO-net, an efficient and scalable method to optimize gradient flow in Wasserstein space. In Section 5, we conclude studying variants of the classical Optimal Transport problem, specifically focusing on AOT by also considering the class of equivalence of filtered processes and associated adapted empirical measures.

2. A continuous idealization of NN for a Mean Field Optimal Control Problem

In this section, we present the basic workflow to treat a feed-forward NN as a dynamical system in order to derive the associated population risk minimization problem.

Indeed, in [35], a continuous idealization of Deep Learning (DL) is introduced, to study the Supervised Learning (SL) problem as am Optimal Control Problem, see, e.g. [27] and [36] for more details.

Following [26], SL aims at estimating the function $\mathcal{F}:\mathcal{X}\to \mathcal{Y}$, commonly known as the Oracle, where $\mathcal{X}$ is a subset of ${\mathbb{R}}^d$ that contains input arrays such as images, texts, and time series, and $\mathcal{Y}$ is the corresponding target set. To achieve this, training begins with a set of K input-target pairs ${\{x_0^i,y_T^i=\mathcal{F}\left(x_T^i\right)\}}_{i=1}^K$ where

• $x_0\in {\mathbb{R}}^d$ denotes the input of the NN;
• $x_T\in {\mathbb{R}}^d$ denotes the output of the NN:
• $y_T\in {\mathbb{R}}^l$ the corresponding target.

For a T-layers network, the Feed-Forward propagation is given by

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

(1)

${\theta }_t$ being the trainable parameters (e.g. bias, weights) in layer t that belong in to a measurable set $\mathcal{U}$ with values in subset of the euclidean space ${\mathbb{R}}^m$.

We aim at minimizing over the set of measurable parameters $\mathcal{U}$ a terminal loss function$\mathsf{\Phi }:\mathcal{X}\times \mathcal{Y}\to \mathbb{R}$ plus a regularization term L over a set of measurable parameters $\mathcal{U}$ to control the size of parameters. In this way, we can state a Supervised Learning problem as an Empirical Risk Minimization problem, namely

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

(2)

2.1. Mean Field Optimal Control Problem

The next steps involve moving from the discrete setting to the corresponding continuous idealization by:

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

Specifically, we pass to the limit both in the layer discretization to obtain a dynamic described by an Ordinary Differential Equation (ODE) instead of finite difference equation (1). Also we identify the input-target pairs as samples from a given distribution ${\mu }_0$ allowing to write a SL problem as a population risk minimization problem.

In summary, the problem aims 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).

More specifically, we consider a probability space $(\mathsf{\Omega },\mathcal{F},\mathbb{P})$ and we assume inputs $x_0$ in ${\mathbb{R}}^d$ being sampled from a distribution $\mathcal{P}\left({\mathbb{R}}^d\right)$ and the corresponding target $y_T$ in ${\mathbb{R}}^l$ belongs to $\mathcal{P}\left({\mathbb{R}}^l\right)$. These variables are considered to be random variables that are jointly distributed according to ${\mu }_0$, which is a distribution in the Wasserstein space ${\mathcal{W}}_2\left({\mathbb{R}}^{(d+l)}\right)$. 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, i.e., measures $\mu$ such that

$${\int }_{X}d{(x,y)}^{p}d\mu \left(x\right)<\infty$$

(3)

for some $y\in X$. Here, $d(x,y)$ is the metric between x and y, and p is a positive real number. The Wasserstein space can be endowed with the Wasserstein distance, also known as the p-th Wasserstein distance, which is a metric that measures the distance between two probability measures $\mu$ and $\nu$ in ${\mathcal{W}}_p\left(X\right)$ as the infimum of the expected cost of transporting the mass from $\mu$ to $\nu$:

$$W_p(\mu ,\nu )={\left(\underset{\gamma \in \mathsf{\Gamma }(\mu ,\nu )}{inf}{\int }_{X\times X}d{(x,y)}^{p}d\gamma (x,y)\right)}^{1/p}$$

(4)

where $\mathsf{\Gamma }(\mu ,\nu )$ is the set of all probability measures on $X\times X$ with marginals $\mu$ and $\nu$.

In our setting, the joint measure ${\mu }_0$ that models the distribution of the input-target pairs and is defined by ${\mu }_0:=\mathcal{P}(x_0,y_T)$ is studied as an element of ${\mathcal{W}}_2$.

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

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

(5)

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

We briefly report basic assumptions we need to have a solution for (5):

• $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 from parameters $\theta$;
• ${\mu }_0$ has finite support in ${\mathcal{W}}_2\left({\mathbb{R}}^{(d+l)}\right)$

Problem (5) 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 [30] and [32] for a viscosity solutions to the HJB equation in the Wasserstein space of probability measures and to [19] for a Pontryagin Maximum Principle for constrained optimal control problems in the Wasserstein space of probability measures.

About some generalizations of the presented result, we cite the article [16] where the authors introduce a BSDE technique to solve the Stochastic Maximum Principle to examine the uncertainty associated with DL. The authors employ an Stochastic Differential Equation (SDE) in place of the ODE appearing in (5) to continuously approximate a Stochastic Neural Network (SNN).

2.2. Empirical measures over controls

For the sake of completeness, we mention also a completely different perspective to describe the learning process in a NN, that deals with the introduction of the empirical distribution of the parameters of neurons (instead of the input-target pairs). Indeed, as illustrated, e.g., by Sirignano and Spilipupouls in [43], it is possible to associate to each layer the corresponding empirical measure and build a measure to describe the whole network that takes values in the space of the called nested measure (also known as adapted Wasserstein distance), that we deepen in Section 5. Following this perspective, the method of Stochastic Gradient Descent (SGD) can be formalized as minimization over probability distribution. Moreover, the training of NN is based on the correspondence between empirical measure of neurons ${\mu }_N$ and the approximated function $f_N$. 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 [23,26,27,29]. Similarly, the convergence to solutions to SDEs have been obtained as universal continuum objects in the small learning rate regime as proved in, e.g., [22].

3. Learning Mean Field Function

In this section, we study some tools in order to learn functions showing a dependency on a mean field term, such as an empirical distribution over a given state space. We start with a brief description of a dynamica of McKean-Vlasov type.

3.1. The McKean-Vlasov SDE

The McKean-Vlasov dynamic involves a SDE that models the behavior of a large population of interacting particles. In this context, the optimal control problem for this system involves finding a control strategy that minimizes a certain cost function while satisfying the dynamics of the system. In particular, the McKean-Vlasov dynamics can be written as:

$$d{X}_t=b\left(X_t\right)dt+\sqrt{ϵ}d{W}_t+\frac{1}{N}\sum _{i=1}^N\int K(X_t-X_i\left(t\right))d{\mu }_i\left(t\right)dt,$$

(6)

where $X_t$ is the state of the system at time t, $b\left(X_t\right)$ is the drift term, $ϵ$ is the diffusion coefficient, $W_t$ is a standard Brownian motion, N is the number of particles in the population, ${\mu }_i\left(t\right)$ is the empirical measure of the particles at time t, and K is a kernel function. Accordingly, the optimal control problem for the McKean-Vlasov dynamics can be written as:

$$J\left(u\right)=\mathbb{E}\left[{\int }_0^{T}L(X_t,u_t)dt+\varphi \left(X_T\right)\right],$$

(7)

subject to the dynamics of the system:

$$d{X}_t=b\left(X_t\right)dt+\sqrt{ϵ}d{W}_t+\frac{1}{N}\sum _{i=1}^N\int K(X_t-X_i\left(t\right))d{\mu }_i\left(t\right)dt+u_{t}dt,$$

(8)

where $u_t$ is the control input, $L(X_t,u_t)$ is the instantaneous cost, and $\varphi \left(X_T\right)$ is the terminal cost.

A possible approach to the optimal control problem consists in using the Pontryagin Maximum Principle, which gives necessary conditions for optimality. The Hamiltonian of the system is defined as:

$$H(X_t,u_t,p_t)=L(X_t,u_t)+p_t·\left[b\left(X_t\right)+\frac{1}{N}\sum _{i=1}^{N}K(X_t-X_i\left(t\right))\right]+\frac{ϵ}{2}{\left|p_t\right|}^2,$$

(9)

where $p_t$ is the adjoint variable. The optimal control is given by:

$$u_t^{*}=arg\underset{u}{min}H(X_t,u,p_t).$$

(10)

The adjoint variable satisfies the following (backward) differential equation:

$$d{p}_t=-\frac{\partial H}{\partial X}(X_t,u_t^{*},p_t)dt,p_T=\frac{\partial \varphi }{\partial X}\left(X_T\right).$$

(11)

3.2. Mean Field Optimal Transport

In the recent work [12], the authors introduce the Mean Field Optimal Transport (MFOT) problem in the field of Mean Field Control, namely optimal control for McKean-Vlasov dynamics. Differently from classical Mean Field Game (MFG) theory, the final distribution is prescribed while all the agents cooperate in order to minimize a total cost without terminal cost. Indeed we follow the numerical scheme introduced in section 3.1 in [12] to approximate feedback controls.

Let ${\mathbb{R}}^d$, describe the state space and denote by ${\mathcal{P}}_2\left({\mathbb{R}}^d\right)$ the set of square-integrable probability measures on ${\mathbb{R}}^d$. Let $f:{\mathbb{R}}^d\times {\mathcal{P}}_2\left({\mathbb{R}}^d\right)\times {\mathbb{R}}^k\to \mathbb{R}$ be the running cost, $g:{\mathbb{R}}^d\times {\mathcal{P}}_2\left({\mathbb{R}}^d\right)\to \mathbb{R}$ be the terminal cost, $b:{\mathbb{R}}^d\times {\mathcal{P}}_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{P}}_2\left({\mathbb{R}}^d\right)$ the aim of MFOT is to compute the optimal feedback control along all the trajectories $v^{*}:[0,T]\times {\mathbb{R}}^d$ 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]$$

(12)

being ${\mu }^v\left(t\right)$ the distribution of $X_t^v$ under the constraint $X^v$ solves the SDE

$$\left\{\begin{array}{c}{X}_0^v\sim {\rho }_0{X}_T^t\sim {\rho }_{T}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.$$

(13)

Always in [12], the authors present different numerical methods to solve MFOT:

• Optimal control via direct approximation of controls v;
• Deep Galerkin Method for solving a forward-backward systems of PDEs;
• Augmented Lagrangian Method with Deep Learning exploiting the variational formulation of MFOT and the primal/dual approach.

In this section, we focus on the direct method (1) to approximate controls of feedback type by an optimal control formulation. The controls are assumed of feedback form and they are approximated by

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

(14)

being $G:{\mathbb{R}}_{+}\to {\mathbb{R}}_{+}$ is an increasing function.

The idea is to use the function in Eq. (14) as a penalty for being far from the target distribution ${\rho }_T$ instead of the classical terminal cost typical of the MFG/MFC literature. Moreover, discretizing the SDE (13) by an Euler-Maruyama scheme as well as restricting the set of controls v to ones approximated via NNS $v_{\theta }$ with parameters $\theta$ helps the development of numerical algorithms. See Section 3.1 in [12] for further details.

3.3. Alternative methods

Another proposed data-driven approach, presented in [8], 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 [37], 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.

4. Learning Gradient Flow

Numerous problems in ML and statistics deal with minimizing some objective function over probability distributions. A powerful tool for optimizing functionals relies on gradient flow in the Wasserstein space.

4.1. Gradient Flow in Wasserstein space

Beyond OT problems, a popular use of Wasserstein distances is to study gradient flows, following the formalism of minimizing movements studied in [3]. This corresponds to discrete implicit stepping (i.e. proximal maps) for the Wasserstein distance (instead of more common Euclidean or Hilbertian metrics) functionals, for instance to model crowd motions.

In the discrete-time setting, gradient flow in Wasserstein space can be used to study the evolution of probability measures over a sequence of discrete time steps. Consider a probability measure ${\rho }^t$ over a finite Euclidean space. The gradient flow of the Wasserstein distance can then be used to find a sequence of probability measures ${\rho }^t$, representing, e.g., the evolution of a distribution of a cloud of points over a finite set $X\in {\mathcal{R}}^n$. The Wasserstein distance between two consecutive measure ${\rho }^t$ and ${\rho }^{t+1}$ is minimized along this sequence of measures, which can be seen as a discrete-time analogue of the continuous-time gradient flow.

The dynamic formulation of OT can be solved by means of the so-called Jordan-Kinderlehrer-Otto (JKO) scheme, which consists in an implicit Euler discretization of a Wasserstein gradient flow

$${\rho }^k=argmi{n}_{\rho \in {\mathcal{P}}_2\left({\mathbb{R}}^n\right)}\left[\mathcal{F}\left(\rho \right)+{\displaystyle \frac{1}{2\tau }}{\mathcal{W}}_2^2({\rho }^{k-1},\rho )\right],$$

(15)

with $\mathcal{F}$ being an energy functional and $\tau$ the time-discretization update. The solution of this consecutive minimization problems result in a sequence of probability measures in $P_2\left({\mathbb{R}}^n\right)$.

Figure 1. Proximal OT for a population dynamics (from [18]).

Figure 1. Proximal OT for a population dynamics (from [18]).

The idea is to move from optimization over probability measure of Eq. (15) to an optimization over convex function, which is still challenging, especially in high dimensions. From Brenier’s theorem, for any $\rho \in P_{2,ac}$, there exists a unique ${\rho }^{k-1}$-measurable gradient $\nabla \psi :{\mathbb{R}}^D\to {\mathbb{R}}^D$ of a convex function $\psi$ s.t. $\rho =\nabla \psi \#{\rho }^{(k-1)}$

4.2. Stochastic Optimization with ICNN and JKOnet

A large variety of dedicated numerical schemes has been proposed for spatial discretization and solving of these time-discrete flows, such as for instance finite differences and Lagrangian schemes. A bottleneck of these schemes is the high computational complexity due to the resolution of an OT-like problem at each time step.

In [7], the authors propose an approach that relies on Input-Convex Neural Networks (ICNN), a class of deep models with a convex output with respect to their inputs introduced in [13], to parameterize the space of convex functions $\left\{{\psi }_{\theta }\right\}$in order to approximate the JKO scheme (15) through

$${\theta }^{*}=argmi{n}_{\theta }\left[\mathcal{F}(\nabla {\psi }_{\theta }\#{\rho }^{(k-1)})+{\displaystyle \frac{1}{2h}}{\int }_{{\mathbb{R}}^D}\left|\right|x-\nabla {\psi }_{\theta }\left(x\right){\left|\right|}_2^{2}d{\rho }^{(k-1)}\left(x\right)\right]$$

The authors in [18] investigate the dynamics of a diverse set of points within the framework of Proximal OT using a Jordan-Kinderlehrer-Otto (JKO) flow of measures. These points are observed periodically at different timestamps. The authors propose a 2-level algorithm to capture the trajectories of this population. Specifically, at time $t+1$, a new configuration is generated that reduces the energy while remaining close to the previous configuration observed at time t, as measured by the Wasserstein distance.

To learn an energy function only through snapshots, the JKOnet algorithm is proposed to compute the JKO flow by a 2-time scale learning algorithm:

- an upper level optimization aiming at minimizing the Wasserstein distance w.r.t. to an energy functional $J_{\psi }$.

- a lower level optimization of the constrained parameters $\theta$ of a ICNN, [13].

In Figure 2, we can see an observed trajectory $\left({\mu }_0,...,{\mu }_T\right)$ of point clouds, we seek parameters $\psi$ for the energy $J_{\psi }$ such that the predictions ${\rho }_1,...,{\rho }_T$ following a JKO flow from ${\rho }_0={\mu }_0$ are close the observed trajectory, by minimizing (as a function of $\psi$) the sum of Wasserstein distances between ${\rho }_{t+1}$, the JKO step from ${\rho }_{t-1}$ using $J_{\psi }$, and data ${\mu }_{t+1}$.

In [31], the authors deal with deep network generative models focusing on the problem of comparing two probability distributions that are degenerate, i.e. distribution in higher-dimensional spaces supported on low-dimensional manifolds, while estimating the parameters of a chosen model that fits observed data.

5. From OT to AOT

After reviewing some variants of the classical OT problem, we present some definitions concerning AOT, i.e. an adapted version of OT.

5.1. OT variants

We start with a brief review of classical OT. 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 reads

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

(16)

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

We focus on the setting where $\mu$ and $\nu$ are distributions computed, e.g., from time series 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)$$

(17)

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

Some generalizations of the classical 2-marginal OT problem concern:

• Martingale Optimal Transport. A first model to encode a temporal adapted structure into OT problems moves within the theory Martingale Optimal Transport (MOT) dealing with transport plans $\pi$ with martingale couplings. The stability of a sequence of transport plans $\{{\pi }^1,...,{\pi }^N\}$ with martingales marginals has been proved to hold in [10], also for the Weak Optimal Transport.
• Multi marginal optimal transport problem. Multi-marginal optimal transport problems introduced in [4] are formulated as follows. Consider $M\in \mathbb{N}$ probability measures ${\mu }_1,...,{\mu }_M$ on ${\mathbb{R}}^d$ and consider the optimization problem

  $$\underset{\pi \in \mathsf{\Pi }({\mu }_1,...,{\mu }_M)}{inf}{\int }_{{\left({\mathbb{R}}^d\right)}^M}c\left(x_1,...,x_M\right)d\pi \left(x_1,....,x_M\right)$$

  (18)

  with c being a lower semi-continuous cost function defined on ${\left({\mathbb{R}}^d\right)}^M$ with marginal laws ${\mu }_1,...,{\mu }_M$ and $\mathsf{\Pi }({\mu }_1,...,{\mu }_M)=\{\pi \in {\mathbb{R}}^{\mathcal{d}}\times ...\times {\mathbb{R}}^{\mathcal{d}}|\forall 1\le i\le M,{\int }_{{\mathbb{R}}^d}d\pi ={\mu }_i\}$
  Trying to approximate this problem by discretizing the state space through N points, $x^1,...,x^N\in {\mathbb{R}}^d$ chosen and a priori fixed, lead to a linear programming problem of size $N^M$.
• From Martingale-constrained to Moment Constrained Optimal Transport . In [5,6], the authors develop algorithms to sample from probability distribution while preserving the convex order enabling to use linear programming solver.
  The same authors study the same problem also in [4] following another perspective by relaxing the martingale constraint into a finite number of moment constraints by means of N real-valued bounded (test) function ${\varphi }_1,...,{\varphi }_N$ defined on ${\mathbb{R}}^d$.
  This method is called Moment Constrained Optimal Transport (MCOT). Differently from multi marginal setting, it is not the state space to be discretized but the marginal laws constrained of Eq. (18) are relaxed into a finite number of moment constraints via some given ${\left\{{\varphi }_N\right\}}_{1,...,N}$ It solves the following:

  $$inf{\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}c\left(x,y\right)d\pi \left(x,y\right)$$

  (19)

  where the infimum is computed over the set of probability measures $\pi$ with values in ${\mathbb{R}}^d\times {\mathbb{R}}^d$ satisfying for all $1\le i,j\le N$

  $${\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}{\varphi }_i\left(x\right)d\pi (x,y)={\int }_{{\mathbb{R}}^d}{\varphi }_i\left(x\right)d\mu \left(x\right)$$

  (20)

  and

  $${\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}{\varphi }_j\left(x\right)d\pi (x,y)={\int }_{{\mathbb{R}}^d}{\varphi }_j\left(y\right)d\nu \left(y\right)$$

  (21)
  Always in [4], they prove that MCOT is a generalization of Martingale constrained OT when by means of $L^1$-valued test functions and characteristic functions to enforce the martingale constraint ( see Equation (3.3.11) from [4]).

5.2. Adapted Optimal Transport

Adapted Optimal Transport (AOT) is the natural problem to deal with laws of stochastic processes, as shown in e.g. [1,14,15,34,38]. Intuitively, causality encodes an adapted (nonanticipative) constraint into the allocation of mass of the classical optimal transport problem. The relationship between causal plans and adapted processes is the adapted analogous of the one between transport plans (Kantorovich) and classical transport maps (Monge).

The basic idea of COT relies on the assumption that the transport maps T introduced in (17)have to be adapted.

If, in the classical setting one can split T component-wise by $T\left(x\right)=(T_1,\left(x\right),...,T_d\left(x\right)$, in the COT setting the adapted maps are defined in the form $T\left(x\right)=(T_1,\left(x_1\right),T_2(x_1,x_2)...,T_d\left(x\right)$ where $T_t$ depends only on the past component of x, namely $x_1,...,x_t$. Thus, the adapted analogous of (17) is defined as

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

(22)

Concerning coupling, the adaptness of T means that under $\pi$, the first component of Y until component t should be independent of the component of X after t. This construction allows to include the conditional independence of Y into the set ${\mathsf{\Pi }}_c$ of arbitrary causal couplings that is defined by

$$\begin{array}{c}{\mathsf{\Pi }}_c(\mu ,\nu )=\{\pi \in \mathsf{\Pi }(\mu ,\nu ):\mathit{if}(X,Y)\sim \pi ,\hfill \\ \mathit{then}(Y_1,...,Y_t)\perp \perp (X_{t+1},...,X_d)\left|(X_1,...,X_t)\right\}\hfill \end{array}$$

Optimizing over the set of causal coupling let us formulate the first $AOT$ problem

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

(23)

that is called Causal OT.

In order to provide a symmetrical property that is desirable to give a more robust structure to the problem and also necessary to introduce a proper notion of adapted Wasserstein distance, one may introduce the set of bicausal couplings ${\mathsf{\Pi }}_{bc}$

$$\begin{array}{c}{\mathsf{\Pi }}_{bc}(\mu ,\nu )=\{\pi \in \mathsf{\Pi }(\mu ,\nu ):\mathit{if}(X,Y)\sim \pi ,\hfill \\ \mathit{then}(Y_1,...,Y_t)\perp \perp (X_{t+1},...,X_d)\left|(X_1,...,X_t)\right\}\hfill \\ \mathit{and}(X_1,...,X_t)\perp \perp (Y_{t+1},...,Y_d)\left|(Y_1,...,Y_t)\right\}\hfill \end{array}$$

(24)

and the equivalent Bicausal-OT formulation

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

(25)

To fix nomenclature, we call Causal OT (23) and Bicausal OT (25), AOT problems. As recently proved in [10], it is possible to write a Linear Program formulation for AOT problems (23) and (25) by discretizing the marginal distributions, see Lemma 3.11 in [10] for further details.

As shown in several works, e.g. [20,24,40], to compute efficiently a solution of OT problems, it is convenient to work with a regularized version of (16) by introducing an entropic regularization in order to apply Sinkhorn’s algorithm [41]. We refer to the landmark book [40] for a complete overview with a detailed treatment of Sinkhorn application in Section 4, [40].

In the same fashion, we provide to Eq. (25) an entropic regularization term to study the following Entropic-AOT

$$\underset{\pi \in {\mathsf{\Pi }}_{bc}(\mu ,\nu )}{inf}\int c(x.y)\pi (dx,dy)+\epsilon D_{KL}(\pi ,\mu \otimes \nu ),$$

(26)

being $\epsilon$ a positive positive constant, $D_{KL}$ the Kullback-Leibler divergence and $\mu \otimes \nu$ the product measures between marginals.

In [10] the authors construct an adapted version of Sinkhorn algorithm allowing to handle computations by constructing a sequence of transport plans ${\pi }^k$ that converges (lineraly) to the optimal value of Eq. (26).

In next paragraphs, we review some of the basic tools of the causal framework: the Adapted Wasserstein distance inherited from AOT, the class of equivalence of filtered processes and the Adapted empirical measure.

5.3. Adapted Wasserstein Distance

The starting point is to consider stochastic processes on the canonical space in finite discrete time with values in ${\mathbb{R}}^d$. We can represent the law of a stochastic process as a probability measure on $\mathcal{P}\left({\mathbb{R}}^d\right)$. Even at this very early stage, it is clear that the usual weak topology on $\mathcal{P}\left({\mathbb{R}}^d\right)$ is not enough for encoding the temporal structure of the process that is crucial for stochastic optimization problem.

Thus, different researchers have introduced topological structure on the set of stochastic processes in order to capture the temporal structure. Among different approaches on defining a distance for multistage stochastic optimization, we focus on nested distance also known as adapted Wasserstein distance.

The notion has been independently defined by Pflug and Pichler [38] in a pure distribution setting, i.e. without reference to probability spaces, by means of particular stochastic processes, called tree processes, and by Lassalle [34] that introduces the so-called causal transport plans in Polish spaces.

Following a dynamical setting, adapted Wasserstein distance can model a distance between laws of stochastic processes integrating also the temporal structure of the data.

Recalling notation and definition from Eq. (25), we define the adapted Wasserstein distance as

$${\mathcal{AW}}_p^p(\mu ,\nu )=\underset{\pi \in {\mathsf{\Pi }}_{bc}(\mu ,\nu )}{inf}{\mathbb{E}}^{\pi }\left[\right||X-{Y\left|\right|}_p^p],p\in [1,+\infty ]$$

(27)

with ${\mathsf{\Pi }}_{bc}$ being the set of bicausal couplings defined in Eq. (24).

The nested distance represents a compatible metric for the weak adapted topology. More precisely, as proven in [9], the convergence ${\mathcal{AW}}_p$ is equivalent to the convergence in the weak topology plus convergence of the $p^{th}$ moment. On the other hand, a possible limitation is due to the fact that $\left({\mathcal{P}}_p\left({\mathbb{R}}^d\right),{\mathcal{AW}}_p\right)$ is not complete since ${\mathcal{P}}_p\left({\mathbb{R}}^d\right)$. According to [11], a possible method to store the extra information needed for achieving the completion of $\left({\mathcal{P}}_p\left({\mathbb{R}}^d\right),{\mathcal{AW}}_p\right)$ is by introducing the following space of filtered stochastic processes.

5.3.1. The class of filtered processes

We recall Definition 1.1 from [11]: a five-tuple

$$\mathbb{X}\left(\mathsf{\Omega },\mathcal{F},\mathbb{P},{\left({\mathcal{F}}_t\right)}_{i=1}^N,{\left(X_t\right)}_{i=1}^N\right),$$

is called a filtered stochastic processes if ${\left({\mathcal{F}}_t\right)}_{i=1}^N$ is adapted to ${\left(X_t\right)}_{i=1}^N$.

It is straightforward to extend the notion of adapted Wasserstein distance defined in (27) to filtered processes by the following

$${\mathcal{AW}}_p^p(\mathbb{X},\mathbb{Y})=\underset{\pi \in {\mathsf{\Pi }}_{bc}(\mathbb{X},\mathbb{Y})}{inf}{\mathbb{E}}^{\pi }\left[\right||X-{Y\left|\right|}_p^p],p\in [1,+\infty ];$$

(28)

We report some properties about the Wasserstein space of stochastic processes:

• $\left({\mathcal{FP}}_p,{\mathcal{AW}}_p\right)$ is a geodesic space (assuming $p>1$) and it is the completion of $\left({\mathcal{P}}_p\left({\mathbb{R}}^d\right),{\mathcal{AW}}_p\right)$. The first property suggests the possibility an interpolation, in the sense of McCann, also of stochastic processes that is not possible if one considers the usual Wasserstein space ${\mathcal{P}}_p\left({\mathbb{R}}^d\right)$;
• two filtered processes $\mathbb{X}$ and $\mathbb{Y}$ are identified if ${\mathcal{AW}}_p(\mathbb{X},\mathbb{Y})=0$. Moreover, this equivalence relation encodes also information about their filtration.
• the space of filtered stochastic processes is defined as a set of equivalence classes as in the case of $L^p$ spaces. An equivalence class of ${\mathcal{FP}}_p$ collects all representatives processes that are identical from a probabilistic perspective.

5.3.2. Adapted Empirical Measure

Another key concept in AOT theory concerns constructing empirical measures of stochastic processes also including a dependence on the flow of time.

We report three possible methods to construct adapted empirical measures:

• in [38], adapted empirical distribution is constructed via convolution of a smooth kernel with the empirical measure on paths;
• in [14,15] the authors project paths on a grid by means of a disjoint union of small cubes projecting the stochastic process into the center of these small cubes;
• in [1] the authors show the convergence of adapted empirical measure in ${\mathbb{R}}^d$, removing the assumption of compact support for the underlying distribution. Moreover, the authors introduce the non-uniform adapted empirical measures by defining a non-uniform grid on ${\mathbb{R}}^{dT}$, dealing with cubes of different sizes, i.e. of lower dimension near the origin, bigger far in the margins.

In all works, they establish convergence of the empirical measure to the true distribution.

6. Conclusions

The present review article provides an all-around and diversified overlook of methods at the overlapping of ML, game theory and probability theory. In doing this, we tried to provide a concrete outlook to possible developments of a challenging research path focused on developing robust algorithms for Adapted optimal transport. Indeed, employing the Sinkhorn algorithm with the addition of the entropic regularization potentially allows to derive a method that can handle the curse of dimensionality. On the other hand, OT problems are sensitive to small perturbations in the input data, which can make the solutions unstable. The application of their convergence results in a causal setting could lead to improved training algorithms and more stable models in the field of causal inference for times series.