# Multisymplectic Formulation of Deep Learning Using Mean–Field Type Control
and Nonlinear Stability of Training Algorithm

Nader T. Ganaba<EMAIL_ADDRESS>
###### Abstract

As it stands, a robust mathematical framework to analyse and study various
topics in deep learning is yet to come to the fore. Nonetheless, viewing deep
learning as a dynamical system allows the use of established theories to
investigate the behaviour of deep neural networks. In order to study the
stability of the training process, in this article, we formulate training of
deep neural networks as a hydrodynamics system, which has a multisymplectic
structure. For that, the deep neural network is modelled using a stochastic
differential equation and, thereby, mean–field type control is used to train
it. The necessary conditions for optimality of the mean–field type control
reduce to a system of Euler–Poincaré equations, which has the a similar
geometric structure to that of compressible fluids. The mean–field type
control is solved numerically using a multisymplectic numerical scheme that
takes advantage of the underlying geometry. Moreover, the numerical scheme,
yields an approximated solution which is also an exact solution of a
hydrodynamics system with a multisymplectic structure and it can be analysed
using backward error analysis. Further, nonlinear stability yields the
condition for selecting the number of hidden layers and the number of nodes
per layer, that makes the training stable while approximating the solution of
a residual neural network with a number of hidden layers approaching infinity.

##  1 Introduction

Machine learning has gained an incredible amount of attention in recent years,
partly thanks to the recent achievement and successes in solving real-world
problems. Problems range from data-driven drug discovery such as using
autoencoders [43] and reinforcement learning to generate new molecules [101],
to others for reducing the dimensionality of single-cell RNA sequencing [120].
Not just computational bioinformatics, but also image recognition [38, 116,
72], and controlling vehicles [127, 39] to name a few. As architectures become
more elaborate, training them becomes more challenging and this has motivated
the influx of techniques from other mathematical disciplines, such as optimal
control [121, 82, 122] and optimal transport [111]. Despite that, there is
still no mathematical framework in which to examine the behaviour of deep
networks, or the convergence of training methods. In [12], optimal control
theory was nominated as a framework for studying the stability of training,
and in [50], it was studied by examining the stability discretisation methods
for ordinary differential equations. As optimal control problems are linked to
fluid dynamics [17], we extrapolate on the ideas of [50, 12] by considering
the optimal control problem from the perspective of multisymplectic
formulation of hydrodynamics. Along with [40], the aim is to translate deep
learning into the language of geometric mechanics in order to gain more
insight into the comportment of deep networks. Profiting from the
multisymplectic nature formulation, in this article, we introduce a training
algorithm based on a multisymplectic numerical scheme. This type of schemes is
the exact solution of another dynamical system whose Lagrangian and
Hamiltonian functions are found using backward error analysis. Knowing the
Hamiltonian for the numerical scheme, we use energy–Casimir method to find the
conditions for the nonlinear stability of the numerical scheme. One of the
benefits of studying the nonlinear stability is that it provides a method to
heuristically select the network’s parameters such number of hidden layers and
nodes per layer.

The route to this formulation starts with the fact that, as pointed out in
[121], residual neural networks, also referred to as _ResNets_ , can be
regarded as the Euler discretisation of a dynamical system. Thus, it is
possible to take the continuous limit and treat it as a neural network with an
infinite number of hidden layers. To train such networks, one possible way is
to use optimal control. The solution of the optimal control problem is
obtained using Pontryagin’s Maximum Principle [106], which states that the
trajectory of the dynamical system solves a Hamiltonian system. Thus,
Pontryagin’s Maximum Principle provides a passage to symplectic geometry [3].

Optimal control training methods use ordinary differential equation for
ResNets [82, 50, 12], here, we use a more general representation that accounts
for uncertainties. Thus, optimal control problem becomes a stochastic one and
that introduces additional challenges. There are many ways to solve stochastic
optimal control problems and here we limit ourselves to only mean–field type
control. The reasoning behind this choice is that mean–field type control
formulates a stochastic control problem as a deterministic one, and in theory.
Mean–field type control originated from mean–field games, which is a
mathematical framework developed to control an ensemble of identical agents
and it was formulated in [58]. At the same time and independently, mean–field
games was introduced in [73, 74, 75] as a framework for modelling differential
games of an infinite number of agents in a population. Its treatment for
optimal control problem was later discussed in [14, 13] and references
therein. For deep learning, it was employed in in [122], where training was
done using mean–field games, although consists of a system of forward and
backward stochastic differential equations, known as McKean–Vlasov dynamics
[26, 25]. On the contrary, the mean–field games system used here consists of a
system of a coupled Fokker–Planck and Hamilton–Jacobi–Bellman equations and
both are partial differential equations. As mentioned, for the deterministic
case, the optimal solution satisfies a Hamiltonian system, which is a
covariant Hamiltonian system in this case. That said, symplectic geometry does
not cover covariant Hamiltonian systems, and for this reason, multisymplectic
geometry is required, which is a generalisation of symplectic geometry [86].

Despite it alluded to in [86], the link between multisymplectic geometry and
Pontryagin’s Maximum Principle is not well explored. When the network
parameters, treated as control, are elements of the Lie algebra associated
with the group of diffeomorphisms, the optimal control problem becomes a
special case known as _Clebsch optimal control_ [41]. This is vital also for
reduction. In [31], it was shown that the stationary variations of the Clebsch
variational principle for hydrodynamic systems yield a multysymplectic system,
which is the a system of three equations: Fokker–Planck,
Hamilton–Jacobi–Bellman and the optimality condition. The Fokker–Planck
equation is a forward partial differential equation, while the
Hamilton–Jacobi–Bellman equation is a backward partial differential equation.
Using the condition that the network’s weights are vector fields of the group
of diffeomorphisms, the mean–field games system is reduced, by eliminating the
Hamilton–Jacobi–Bellman equation, to a system of forward equations. The result
is a system of hydrodynamic equations known as _Euler–Poincaré equation with
advected density_[56]. Many hydrodynamic equations, such as the Kortweg-de
Vries equation [102], the Camassa-Holm equation [23, 69] are part of the
Euler–Poincaré family of equations. From the calculus of variations, the
trajectory that minimises a cost functional satisfies the Euler-Lagrange
equation, and this case is the mean–field system, and when there is symmetry,
the variational problem is reduced and the optimal trajectory satisfies an
equivalent the Euler–Poincaré equation. With that, the claim here is that the
weights of the deep residual neural network satisfy a variant of the
Euler–Poincaré equation with advected quantity on the group of
diffeomorphisms.

As the nature of data is discrete, the optimal control problem needs to be
projected to a finite-dimensional space by using numerical schemes. For
systems whose dynamics are described by an ordinary differential equation, the
Hamiltonian equations of motion for the optimal control problems have a
symplectic structure, hence geometric integrators can be used to solve optimal
control problems numerically as it was done [64, 99, 33]. The advantage of
using geometric integrators is their stability and preservation of the
underlying geometry, thus eliminating the possibility of the emergence of
unrealistic solutions. Due to the fact that mean-field type optimal control is
multisymplectic in nature, we thus use multisymplectic integrators derived
using a discrete variational method [86], as an extension to discrete
mechanics optimal control of [64, 99, 33]. This projection is important, as it
provides a systematic way to train a network with a finite number of hidden
layers, and concurrently simulates the behaviour of a network with an infinite
number of hidden layers. To verify the presented training algorithm using a
multisymplectic variational integrator, two toy problems are solved: a density
matching problem from variational inference and images of handwritten digits
generation of using a variational autoencoder. It is worth pointing out that a
multisymplectic integrator can be applied to problems in machine learning in a
different way as in [9], where the data points are elements of a Lie group.
Whereas, here, the network parameters belong to a Lie algebra, and the
multisymplectic structure arises from mean–field type control for stochastic
optimal control. The data need not to belong to a Lie group to use our
approach to multisymplectic integrators.

Reducing the optimal control problem to a hydrodynamics problem and solving it
using a multisymplectic scheme allows us to employ the same methods used to
analyse the stability of fluids [7, 6, 55, 54] to determine the stability of
the training algorithm. An intermediate step is required to carry out
stability analysis is backward error analysis. In short, backward error
analysis seeks to find an equation, known as _the modified equation_ , whose
exact solution is the numerical solution obtained by using the multisymplectic
integrator. In [61], it was shown that multisymplectic integrators do have a
multisymplectic modified equations. For our analysis, the modified equation
helps to derive the error equation for the numerical error and show that it is
nonlinearly stable. This means if the error is small enough, then it remains
bounded. To establish nonlinear stability of a system, we use a method called
_energy–Casimir method_ [55]. It is an algorithmic approach to constructing a
norm for determining if the equilibrium solution is nonlinearly stable. This
is done by adding a conserved quantity known as the Casimir function to the
Hamiltonian. Network parameters such as the number of hidden layers, number of
nodes and learning rate appear in the modified equation, and nonlinear
stability provides conditions for the said parameters to ensure that the
training problem converges. The main difference between the approach taken
here and [12], is that the latter is based on exponential stability and uses
Lyapunov’s notion of stability.

_The outline of this paper:_ Section 2 gives a brief review of topics, and
notions that are used throughout this article. The recap starts with giving
the most important definitions of multisymplectic geometry, in section 2.1,
and then moves to optimal control theory in section 2.2. We also give a
statement of the Pontryagin’s Maximum Principle, without providing the proof
as it is too technical to review here. The review is concluded with a
thumbnail sketch of mean–field type control 2.3. In section 3, a formulation
of deep learning problem as a Fokker–Planck Hamilton–Jacobi–Bellman mean–field
games system is presented and such problem is shown to be a multisymplectic
system. The mean–field games system is reduced to obtain a Euler–Poincaré
equation with advected parameter. A numerical scheme for the multisymplectic
system for deep learning is presented in section 4 and it is based on
multisymplectic variational integrator and it is demonstrated how it can be
used for practical problems: density matching, and variational autoencoder,
discussed in section 4.1.1 and 4.1.2, respectively. Backward error analysis
for the multisymplectic scheme is carried out in section 5. Section 6 concerns
with the study of the nonlinear stability of multisymplectic scheme, where the
conserved quantities and Noether’s theorem for deep learning are discussed in
section 6. A Hamiltonian formulation is presented in section 6.1 and it is
need to apply the energy–Casimir method, which is reviewed in section 6.2. The
main result of nonlinear stability is presented in section 6.3 along with
conditions on the number of nodes and hidden layers that guarantee stable
training.

##  2 Preliminaries

###  2.1 Review of multisymplectic geometry

In this section, we provide a brief description of multisymplectic geometry to
familiarise the reader with the notation and basic concepts used throughout
this paper. Most of the definitions and concepts mentioned here are discussed
in detail in [86, 89]. The core of the numerical algorithm of section 4 is
rooted in multisymplectic geometry and the idea is to discretise the infinite-
dimensional equations to finite-dimensional ones in a way that retains most of
the geometry. Although this adds an extra level of complexity, it is important
in understanding the behaviour of the algorithm.

The interest regarding multisymplectic geometry stemmed from the attempts at
finding a formulation that would generalise classical Hamiltonian dynamics to
account for covariant field-theoretic Hamiltonian [34, 24, 47, 37, 48]. One of
the first encouraging examples of the promises of multisymplectic formulation
was the recovery of the symplectic structure of the Korteweg-de Vries
equation, which has a covariant Hamiltonian density [46]. The generalisation
to a larger family of Hamiltonian partial differential equations, written in
an abstract form,

$\displaystyle M\partial_{t}z+K\partial_{x}z=\partial_{z}S(z),$ (2.1)

was introduced in [19]. A Lagrangian equivalent to the multisymplectic
formulation based on variational principle was presented in [86]. Most
importantly, the aforementioned paper introduced variational integrators for
partial differential equations. The link between both Hamiltonian and
Lagrangian versions of multisymplectic theory was explored in detail in [89].

Let $Q$ be a smooth manifold. In Hamiltonian mechanics, the trajectories
belong to the cotangent bundle, $T^{\ast}Q$, and for configuration manifold,
as opposed to the tangent bundle, $TQ$, in Lagrangian mechanics. A curve on
$Q$ is parameterised by an independent variable, in many cases the independent
variable is time $t$, and without the loss of generality, let $t\in B$, where
$B\subset\mathbb{R}$, be the base manifold. The image generated by $q(t)$
defines a subset of manifold $Q$. We also have the preimage $\pi:Q\to B$. When
the curve $q$ depends on more than one variable, the base manifold $B$ is an
$n+1$ dimensional manifold, with $n$ spatial dimensions and one temporal
dimension. For example, when we have a $(1+1)$ partial differential equation
the base manifold is chosen to be
$B:=[l_{0},l_{1}]\times[t_{0},t_{1}]\subset\mathbb{R}^{2}$. The configuration
manifold is instead a fibre bundle $\pi_{BE}:E\to B$. The preimage
$\pi^{-1}_{BE}:B\to E$ of $E$ over a point $x\in B$ denotes the _fibre at
point $x$_. A continuous function $\phi:B\to E$ is the _section_ of $\pi_{BE}$
and we use $\Gamma(E)$ to denote all the section on $E$. In the language of
partial differential equations, a section in $\Gamma\left(\pi_{BE}\right)$
corresponds to the solution of the partial differential equation.

The rate of change of the curve $q(t)\in Q$ with respect to $t$, denoted by
$\dot{q}(t)$, is the derivative of $q(t)$ and it belongs to the tangent bundle
$TQ$, and the covariant analogue of that is _the first jet bundle_ $J^{1}(E)$.
It is defined as a bundle over $E$, whose fibres are the mapping
$\gamma:T_{x}B\to T_{y}E$ such that
$T_{\pi_{BE}}\circ\gamma=\mathsf{Id}_{T_{x}B}$. To put it into other words, a
first jet bundle is an equivalence relation between two curves evaluated at
the same point that are equal up to their first-order derivatives. Using the
same construction, once can construct higher jet bundles, the $(n+1)^{th}$ jet
bundle is a fibre bundle on the $(n)^{th}$ jet bundle.

Figure 1: A pictorial representation of a fibre bundle $\pi_{BE}:E\to B$.

As mentioned before, the base manifold $B$ is an $n+1$ dimensional manifold,
and its coordinates are $x^{\mu}$ where $\mu=1,\dots,n,0$ and the fibres on
$E$ have coordinates $y^{A}$, with $A=1,\dots,N$. Meanwhile, the first jet
bundle $J^{1}(E)$’s coordinates are denoted by $v_{\mu}^{A}$. Now, we are
interested in sections on $J^{1}(E)$ and specifically a section for $B\to
J^{1}(E)$ and to do so, we define the section $\phi:B\to E$ and to “lift it”
to $J^{1}(E)$, we take its tangent map $T_{x}\phi$. The section on $J^{1}(E)$
over $B$ is denoted by $j^{1}(\phi)$ and it is known as the _first jet
prolongation_ and in coordinates it is given by

$\displaystyle x^{\mu}\mapsto\left(x^{\mu},\phi^{A},\frac{\partial}{\partial
x^{\mu}}\phi^{A}\right).$

Cotangent bundles are regarded as the dual of the tangent bundle. The same is
true here, _the dual of the first jet bundle_ $J^{1}(E)^{\ast}$ is the
covariant analogue of the cotangent bundle and it is the space of maps of
$J^{1}(E)$ to $\Lambda^{n+1}(B)$, the bundle of $n+1$-forms over $B$. The
coordinates of $J^{1}(E)^{\ast}$ are given by

$\displaystyle(p,p_{A}^{\mu}),\quad
v_{\mu}^{A}\to(p+p_{A}^{\mu}v_{\mu}^{A})d^{n+1}x,$

and here the differential form $d^{n+1}x$ is $dx^{1}\wedge
dx^{2}\wedge\dots\wedge dx^{n}\wedge dt$. Next quantities that are important
to the multisymplectic formulation are the $n+1$ and $n+2$ forms that are
covariant counterparts of the canonical one and two-forms in symplectic
geometry. These forms play a pivotal role in multisymplectic geometry as we
shall see later. We first define the projection
$\pi_{E(J^{1}(E)^{\ast})}:J^{1}(E)^{\ast}\to E$ is the projection from the
bundle of $\Lambda^{n+1}$ forms to $E$, then we define the sub-bundle $Z$ by

$\displaystyle
Z:=\left\\{z\in\Lambda^{n+1}_{y}\,|\,V\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}(W\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}z)=0,\quad
V,W,\in V_{y}E\right\\},$

where $V_{y}E$ is the vertical sub-bundle of $TE$ defined by

$\displaystyle V_{y}E:=\left\\{W\in T_{y}E\,|\,T\pi_{BE}\cdot W=0\right\\},$

and it can be described as the space of elements of $T_{y}E$, which renders
the derivative of $\pi_{BE}$ along their direction to be zero. The elements of
$Z$ are defined as

$\displaystyle
z=pd^{n+1}x+p_{A}^{\mu}\mathrm{d}y^{A}\wedge\left(\partial_{\mu}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n+1}x\right).$

The canonical $n+1$ form $\Theta_{\Lambda}$ on $\Lambda^{n+1}$ is defined by

$\displaystyle\Theta_{\Lambda}(z)\left(u_{1},\dots,u_{n+1}\right)=\left(\pi_{E\Lambda}^{\ast}z\right)\left(u_{1},\dots,u_{n+1}\right),$

where $\pi_{E\Lambda}:J^{1}(E)^{\ast}\to E$ the projection from the space of
$n+1$ forms to $E$ and $u_{1},\dots,u_{n+1}\in T_{\Lambda}\Lambda^{n+1}$ and
in coordinates that is

$\displaystyle\Theta_{\Lambda}=p_{A}^{\mu}dy^{A}\wedge\left(\partial_{\mu}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n+1}x\right)+pd^{n+1}x.$

The canonical $n+2$ form $\Omega$ is defined by $\Omega=-d\Theta_{\Lambda}$
and in coordinates it is given by

$\displaystyle\Omega=dy^{A}\wedge
dp_{A}^{\mu}\wedge\left(\partial_{\mu}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n+1}x\right)-dp\wedge
d^{n+1}x.$ (2.2)

The Lagrangian density $\mathcal{L}:J^{1}(E)\to\Lambda^{n+1}(B)$

$\displaystyle\mathcal{L}=L(j^{1}(\phi))d^{n+1}x,$ (2.3)

and the Legendre transformation is the map $\mathbb{F}\mathcal{L}:J^{1}(E)\to
J^{1}(E)^{\ast}$ and it is defined by
$\mathbb{F}\mathcal{L}(\gamma)\cdot\gamma^{\prime}=\mathcal{L}(\gamma)+\left.\frac{d}{d\epsilon}\right|_{\epsilon=0}\left(\mathcal{L}(\gamma+\epsilon(\gamma^{\prime}-\gamma))\right)$,
where $\gamma\in J^{1}(E)$. This is used to define the _Cartan form_ , which
is the result of the variational principle. The Cartan form is known as the
Lagrange one-form in the finite dimensional case. What is different regarding
the Cartan form, $\Theta_{\mathcal{L}}$ known as _the multisymplectic form_ ,
is that it an $n+1$ form defined on $J^{1}(E)$ instead of $J^{1}(E)^{\ast}$

$\displaystyle\Theta_{\mathcal{L}}=\left(\mathbb{F}\mathcal{L}\right)^{\ast}\Theta,$
(2.4)

and in coordinates it is expressed explicitly by

$\displaystyle\Theta_{\mathcal{L}}=\left(\frac{\partial
L}{\partial\phi_{\mu}^{A}}\right)d\phi_{\mu}^{A}\wedge\left(\partial_{\mu}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n+1}x\right)+\left(L-\frac{\partial
L}{\partial\phi_{\mu}^{A}}\phi_{\mu}^{A}\right)d^{n+1}x$ (2.5)

Similarly, the $n+2$-from $\Omega_{\mathcal{L}}$ on $J^{1}(E)$ is defined by

$\displaystyle\Omega_{\mathcal{L}}=-d\Theta_{\mathcal{L}}=\left(\mathbb{F}\mathcal{L}\right)^{\ast}\Omega.$

Using the Lagrangian density (2.3), the action functional is defined as

$\displaystyle\mathcal{S}(\phi)=\int\mathcal{L}(j^{1}(\phi\circ\phi_{B}^{-1})),$
(2.6)

where $\phi:U\to E$, $U$ is a smooth manifold with smooth closed boundaries
and $\phi_{B}=\pi_{BE}(\phi):U\to B$. Let
$\Phi:G\times\mathcal{F}\to\mathcal{F}$, where $\mathcal{F}$ is the set of
smooth maps, defined by

$\displaystyle\Phi(\eta_{E},\phi)=\eta_{E}\circ\phi.$ (2.7)

The map $\eta_{E}:E\to E$ is a diffeomorphism and analogously, on the base
space $\eta_{B}:B\to B$. The purpose of using such diffeomorphism is that we
can construct a family of sections on $E$. In this particular case, let
$\eta^{\epsilon}_{E}$ be a smooth curve in $G$ parameterised by $\epsilon$ and
$\eta^{0}_{E}=e$ the identity of $G$. Also, let

$\displaystyle
V=\left.\frac{d}{d\epsilon}\right|_{\epsilon=0}\Phi(\eta^{\epsilon}_{E},\phi),\quad\text{and}\quad
V_{B}=\left.\frac{d}{d\epsilon}\right|_{\epsilon=0}\eta^{\epsilon}_{B}\circ\phi,$
(2.8)

where $V$ is a vector field for the transformation group on $E$ and $V_{B}$ is
a vector field for the transformation group on $B$. With that, the variations
of the action functional (2.6) is

$\displaystyle\delta\mathcal{S}(\phi)=\left.\frac{d}{d\epsilon}\right|_{\epsilon=0}\mathcal{S}(\Phi(\eta^{\epsilon}_{E},\phi))=\left.\frac{d}{d\epsilon}\right|_{\epsilon=0}\int\mathcal{L}(j^{1}(\Phi(\eta^{\epsilon}_{E},\phi)\circ\phi_{B}^{-1})),$
(2.9)

and using the property of diffeomorphisms $\eta^{\epsilon}_{E}$ and
$\eta^{\epsilon}_{B}$ on sections

$\displaystyle
j^{1}(\Phi(\eta^{\epsilon}_{E},\phi)\circ\phi_{B}^{-1})=j^{1}(\eta_{E}\circ\phi\circ\phi_{B}^{-1}\circ\eta_{B}^{-1})$
$\displaystyle=j^{1}(\Phi(\eta^{\epsilon}_{E},\phi)\circ\phi_{B}^{-1})$
$\displaystyle=j^{1}(\eta_{E})\circ
j^{1}(\phi\circ\phi_{B}^{-1})\circ\eta_{B}^{-1}$

and then the variations of the action along the vector field $V$ become

$\displaystyle d\mathcal{S}(\phi)\cdot
V=\int_{U_{B}}j^{1}(\phi\circ\phi_{B}^{-1})^{\ast}\pounds_{j^{1}(V)}\Theta_{\mathcal{L}},$

and using the Cartan’s formula definition of Lie derivative [86], we arrive at

$\displaystyle d\mathcal{S}(\phi)\cdot
V=-\int_{U_{B}}j^{1}(\phi\circ\phi_{B}^{-1})^{\ast}\left[j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Omega_{\mathcal{L}}\right]+\int_{\partial
U_{B}}j^{1}(\phi\circ\phi_{B}^{-1})^{\ast}\left[j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Theta_{\mathcal{L}}\right]$
(2.10)

Setting the $d\mathcal{S}(\phi)\cdot V=0$, the first integral is zero when
$\phi$ satisfies the Euler-Lagrange equation

$\displaystyle
j^{1}(\phi\circ\phi_{B}^{-1})^{\ast}\left[j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Omega_{\mathcal{L}}\right]=0,$
(2.11) or $\displaystyle\frac{\partial L}{\partial
y^{A}}(j^{1}(\phi\circ\phi_{B}^{-1})-\frac{\partial}{\partial
x^{\mu}}\left(\frac{\partial L}{\partial
v_{\mu}^{A}}(j^{1}(\phi\circ\phi_{B}^{-1}))\right)=0.$

When $\phi$ satisfies (2.11) and $V$ and $W$ are two vector fields generated
by the one-parameter transformation group of $\phi$ and the boundary integral
in (2.10) gives rise to the following boundary integral

$\displaystyle\int_{\partial
U_{B}}j^{1}(\phi\circ\phi_{B}^{-1})^{\ast}\left[j^{1}(W)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Omega_{\mathcal{L}}\right]=0,$
(2.12)

and it is known as _the multisymplectic form formula_. The derivation of this
formula can be found in [86]. The importance of it is that it shows that the
solution of the Euler-Lagrange equation (2.11) preserves the multisymplectic
form (2.4) [86].

###  2.2 Review of optimal control

As optimal control is central to this work, the basic definitions are stated
here in order to familiarise the reader with the notation that will be used
throughout this paper. We follow the notation used in [22]. However the topic
is not covered here in great detail and the reader is advised to consult [115,
2, 65] and references therein for detailed exposition.

Given the quadruple $\Sigma=(Q,f,\mathcal{U})$ where $Q$ denotes the
configuration manifold, $f:Q\times\mathcal{U}\to TQ$ denotes the vector field
that describes the motion of the system, and $\mathcal{U}$ denotes the space
of admissible controls. We say the quadruple $\Sigma=(Q,f,\mathcal{U})$ is a
_controlled stochastic system_. The stochastic differential equation that
governs the controlled system $\Sigma$ is the following

$d_{t}q=f(q,\theta)\,dt+\nu\sum_{k}\xi_{k}(q)\circ dW_{t}^{k},$ (2.13)

where $q\in Q$ and $\xi_{k}$ are vector fields that couple the independent and
identically distributed Brownian motions $W^{k}(t)$ to the tangent bundle
$TQ$. Before giving a more rigorous definition of optimal control, we first
need to define the main quantities of interest: the cost Lagrangian and cost
functional. For the stochastic control system $\Sigma=(Q,f,\mathcal{U})$, we
say a continuous, convex map $\ell:\mathcal{U}\times Q\mapsto\mathbb{R}$,
integrable in the time interval $[t_{0},t_{1}]$, is the _cost Lagrangian_. In
some sources, it is referred to as the _cost function_ and that is to avoid
confusion with other Lagrangians. While the Lagrangian, $\ell$, gives the cost
of the state and control at time $t$, but for the cost in the time interval
$[t_{0},t_{1}]$, a functional is used instead. _The cost functional_
$\mathcal{S}_{\Sigma}:\mathcal{U}\times Q\mapsto\mathbb{R}$ is used as a
performance metric and it is defined as

$\mathcal{S}_{\Sigma}=\int_{t_{0}}^{t_{1}}\ell(q(t),\theta(t))dt.$

With that, the definition of a optimal control problem can be stated as
follows:

###### Definition 2.1 (Optimal control problem).

Given a stochastic control system $\Sigma=(Q,f,\mathcal{U})$, let $Q_{0}$ and
$Q_{T}$ be two distinct submanifolds of $Q$, then find the pair
$(q^{\ast},\theta^{\ast})$, where $q^{\ast}\in Q$ and
$\theta^{\ast}\in\mathcal{U}$, such that the curve $q^{\ast}(0)\in Q_{0}$ and
$q^{\ast}(T)\in Q_{T}$ and in the same time satisfying
$\mathcal{S}_{\Sigma}(q^{\ast},\theta^{\ast})<\mathcal{S}_{\Sigma}(q,\theta)$
for every $q\in Q$ and $\theta\in\mathcal{U}$.

Due to the fact that it is indispensable, we only sketch Pontryagin’s Maximum
Principle. As the proof is technical and out of the scope of this paper,
nevertheless, it can be found in [106, 2, 65, 78] and the supplement chapter
[22] of [21]. One of the nuclei of Pontryagin’s Maximum Principle is the
Hamiltonian function for $\Sigma$. For the control system
$\Sigma=(Q,f,\mathcal{U})$ and the cost Lagrangian $\ell$, the _Pontryagin’s
Hamiltonian function_ , also known as _control Hamiltonian_ ,
$H_{\Sigma}:T^{\ast}Q\times\mathcal{U}\mapsto\mathbb{R}$, associated to the
control system $\Sigma$ is defined as

$\displaystyle
H_{\Sigma}(q,\lambda,\theta)=\left\langle\lambda,f(q,\theta)\right\rangle_{T^{\ast}Q\times
TQ}-\ell(q,\theta).$

Once the substituting the expression for optimal $\theta$ in terms of
$(q,\lambda)$, we arrive at _maximum Hamiltonian_ ,
$H^{\ast}_{\Sigma}:T^{\ast}Q\mapsto\mathbb{R}$, and it is defined as

$\displaystyle
H^{\ast}_{\Sigma}(q,\lambda)=\sup\left\\{H_{\Sigma}(q,\lambda,\theta)\quad|\quad\theta\in\mathcal{U}\right\\}.$
(2.14)

###### Theorem 2.2 (Pontryagin’s Maximum Principle [106]).

Let $\Sigma=(Q,f,\mathcal{U})$ be the control system, and let $Q_{0}$ and
$Q_{T}$ be the disjoint sub-manifolds of $Q$. If the pair $(q,\theta)$ solves
the optimal control problem, i.e. minimise the cost functional

$\mathcal{S}=\int_{0}^{T}\ell(q(t),\theta(t))dt$ (2.15)

subject to $\dot{q}(t)=f(q(t),\theta(t))$, with $q(0)=q_{0}\in Q_{0}$ and
$q(T)=q_{T}\in Q_{T}$, then there exists $\lambda_{0}=\\{0,1\\}$ and a
cotangent vector

$\lambda:[0,T]\mapsto T^{\ast}Q$

along $q$ and the Hamiltonian
$H^{\ast}_{\Sigma}(\lambda(t))=\sup\limits_{\theta\in\mathcal{U}}H_{\Sigma}(\lambda(t))$
and $\lambda(t)$ evolves in time according to

$\dot{\lambda}=X_{H^{\ast}_{\Sigma}}(\lambda(t))$

In other words, $\lambda$ is an integral curve of a Hamiltonian vector field
associated with $H^{\ast}_{\Sigma}$. Along with that, it also satisfies the
following transversality condition:
$\lambda(0)\in\mathrm{ann}(T_{q(t_{0})}q_{0})$ and
$\lambda(1)\in\mathrm{ann}(T_{q(t_{1})}q_{T})$ and we either have
$\lambda_{0}=1$ or $\lambda(0)\neq 0$ and finally there exists a constant
$C\in\mathbb{R}_{+}$ such that the maximum Hamiltonian is constant, i.e.
${H^{\ast}_{\Sigma}}(\lambda(t))=C$.

For a specific setting, when the manifold $Q$ is acted by a Lie group action,
the optimality condition yields a momentum map and Hamilton’s equations of the
Pontryagin’s Hamiltonian system, generated by (2.14), can be reduced to the
Euler–Poincaré equation. As a result the initial condition $q_{0}$ and $q_{T}$
both lie to the same $G$-orbit. This observation was pointed out in [41, 16]
and they formulated optimal control in terms of Clebsch variables.

###### Definition 2.3 (Clebsch optimal control [41]).

Let $\Sigma=(Q,f,\mathcal{U})$ be the control system, then the Clebsch
variational principle for optimal control is

$\displaystyle\delta\mathcal{S}=\delta\int_{0}^{T}\ell(q(t),\theta(t))\,dt+\int_{0}^{T}\int_{Q}\lambda\cdot\left(\dot{q}(t)-f(q(t),\theta(t))\right)\,d^{n}x\,dt=0,$

where $\lambda$ is the Lagrange multiplier that acts as a constraint to the
probability density function $\rho$ to evolve according to

$\displaystyle\dot{q}(t)=f(q(t),\theta(t)),$

with $q(0)=q_{0}\in Q_{0}$ and $q(T)=q_{T}\in Q_{T}$, where
$Q_{0},Q_{T}\subset Q$.

###  2.3 Mean–field type optimal control

Here, we briefly overview the mean–field type control of [14, 13]. Consider a
population of $N$ identical agents on $Q$ and each agent has control over the
stochastic differential equation

$dX_{t}^{i}=f(X_{t}^{i},\theta(t))dt+\nu\sum_{k}\xi_{k}(X_{t}^{i})dW^{k}_{t},\quad
X^{i}_{0}=x^{i}\in Q\subset\mathbb{R}^{d},$ (2.16)

where $X_{t}^{i}$ are Markov processes and the drift is denoted by
$\theta(x,t)$, and $W^{i}(t)$ are independent and identically distributed
Brownian motions. Here, $\xi_{k}$ are diffusion vector fields and they are
used to couple $W^{i}(t)$ to the tangent bundle. For mean–field type control,
the goal is to find $\theta=(u(x,t),b(x,t))$ such that it steers the whole
population from an initial $X^{i}_{0}=x_{0}$ to $X^{i}_{T}=x_{T}$. The
solution of mean–field games, as the size of the population of agents goes to
infinity, $N\to\infty$, converges to Nash equilibrium [74].

###### Remark 2.4.

In the previous sections, curves on $Q$ were denoted by $q(t)$. Here, $X_{t}$
is used for stochastic processes on $Q$ and to make the clear distinction
between them and deterministic curves.

###### Remark 2.5.

It is important to point out that when $Q$ is not $\mathbb{R}^{d}$, the
stochastic differential equation (2.16) is replaced by a Stratonovich
stochastic differential equation. When dynamics are manifolds, Itô stochastic
differential equations require extra construction to define them on smooth
manifolds, hence it is practical to use Stratonovich stochastic differential
equation [76, 52]. In this paper, we consider $Q\subset\mathbb{R}^{d}$, in
order to use Itô stochastic differential equation (2.16).

The idea of [73] and [58], instead of considering individual agents, their
states are substituted by a probability density function. Thus, (2.16) is
replaced by the Fokker–Planck equation, which is a partial differential
equation that describes the time evolution of the probability density function
associated a Markov process $X_{t}$, and it is given by

$\displaystyle\frac{\partial}{\partial
t}\rho(x,t)\,d^{n}x+\pounds_{f}{(\rho(x,t)\,d^{n}x)}-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{j}}{\pounds_{\xi_{i}}{(\rho(x,t)\,d^{n}x)}}=0.$
(2.17)

The term $\pounds_{f}$ is the Lie derivative along the vector field $f(x,u)$
and it is defined by

$\pounds_{f}{(\rho\,d^{n}x)}=\left.\frac{d}{ds}\right|_{s=0}(\psi(s))^{\ast}(\rho\,d^{n}x),$

where $\psi^{\ast}$ is the pullback along the diffeomorphism
$\psi:I\times\mathrm{Den}(Q)\to\mathrm{Den}(Q)$ and it is generated by the
vector flow of $f(x,u)$. Since $\rho$ is a density, the Lie derivative of
along $f$ in general coordinates, is written as

$\displaystyle\pounds_{f}{(\rho\,d^{n}x)}=\frac{1}{\sqrt{|\mathbb{G}|}}\partial_{i}\left(\sqrt{|\mathbb{G}|}\rho
f^{i}\right)\,d^{n}x,$

where $\mathbb{G}$ is the space’s metric tensor. In coordinates on
$\mathbb{R}^{N}$ the Lie derivative is

$\pounds_{f}{(\rho\,d^{n}x)}=\mathrm{div}\left(\rho f\right)\,d^{n}x.$

###### Remark 2.6.

The Fokker–Planck equation is usually written in terms of coordinates on
$\mathbb{R}^{N}$, but in (2.17), it is written in terms of Lie derivatives to
stress that fact that we are using a coordinate free description. It is
trivial to arrive at such description. Tout court, one computes the integral

$\displaystyle\int_{Q}f(x)\frac{\partial\rho(x,t|x_{0},t_{0})}{\partial
t}d^{n}x,$

where $f$ is a scalar function $f:Q\to\mathbb{R}$, by using the definition of
a time derivative as the limit of a difference between $\rho$ at time
$t+\delta t$ and $t$ and take the limit as $\delta t\to 0$. The density at
$t+\delta t$ obeys the Chapman-Kolmogorov equation

$\rho(x,t+\delta t|x_{0},t_{0})=\int_{Q}\rho(x,t+\delta
t|y,t)\rho(y,t|x_{0},t_{0})d^{n}y$ (2.18)

and then using writes Itô’s lemma [62, 100] for a scalar function
$g:Q\to\mathbb{R}$ in terms of Lie derivatives

$\displaystyle
g(x)=g(y)+\pounds_{dX^{i}}g(y)+\frac{1}{2}\pounds_{dX^{i}}\pounds_{dX^{i}}g(y),$

we then arrive at the Fokker–Planck equation (2.17).

Going back to the stochastic optimal control problem, generally, the cost
Lagrangian is defined as an expected value

$\ell(\theta,\rho)=\int_{Q}\rho(x,t)\widehat{\ell}(\theta)\,d^{n}x.$ (2.19)

In this note, we restrict ourselves to the case when the
$\widehat{\ell}(\theta)$ is a quadratic function, and this case
$\ell:\mathcal{U}\times\mathbb{R}\mapsto\mathbb{R}$ becomes

$\ell(\theta,\rho)=\frac{1}{2}\int_{Q}\rho(x,t)\left\|\theta\right\|^{2}\,d^{n}x$
(2.20)

We now formulate a method for solving stochastic geometric optimal control
problems. Here, the Fokker–Planck equation (2.17) is a constraint, in lieu of
(2.13). In this setting, the problem is no longer on the state manifold $Q$,
instead it is on the space of probability densities of $Q$, denoted by
$\mathrm{Den}(Q)$. The space of probability density functions on $Q$ is
defined as

$\displaystyle\mathrm{Den}(Q)=\left\\{\rho\in
H^{1}(Q)\quad|\quad\rho>0,\int_{Q}\rho\,d^{n}x=1\right\\}.$ (2.21)

It is worth noting that densities are usually defined as volume forms, but
since the probability densities are related to the diffusion processes we
treat them as functions with a reference volume form. Whichever $Q$ being
finite or infinite-dimensional, $\mathrm{Den}(Q)$ is considered as an
infinite-dimensional manifold for all cases [66].

The space of admissible controls $\mathcal{U}$ remains the same, however the
input data of the training and testing sets are now elements of
$\mathrm{Den}(Q)$ [66]. It is worth noting that the setting used here is the
same for computational fluid mechanics formulation of optimal transport [10].
However, in [10] the equation for the probability density function is
continuity equation, meaning there are no double Lie derivatives as in (2.17)
and control vector field $\theta$ is a potential vector field. Then the
control problem for a stochastic system can be stated as the following
quadruple $\Sigma=(\mathrm{Den}(Q),f,\mathcal{U})$ where $\mathrm{Den}(Q)$
denotes the configuration manifold, $f(x,\theta):Q\times\mathcal{U}\to T(Q)$
denotes the vector field that describes the motion of the system,
$\mathcal{U}$ denotes the space of admissible controls. The dynamics of the
control system are governed by the Fokker–Planck equation

$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f}{(\rho\,d^{n}x)}-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}(\rho\,d^{n}x)=0,$
(2.22)

where $\rho(x,t)\,d^{n}x\in\mathrm{Den}(Q)$ is the probability density
function.

###### Definition 2.7 (Mean–field type optimal control problem [58]).

Given a stochastic control system $\Sigma=(\mathrm{Den}(Q),f,\mathcal{U})$,
let $Q_{0}$ and $Q_{T}$ be two distinct submanifolds of $Q$, then find the
pair $(\rho^{\ast},\theta^{\ast})$, where $\rho^{\ast}\in\mathrm{Den}(Q)$ and
$\theta^{\ast}\in\mathcal{U}$, such that the curve
$\rho^{\ast}(x,0)\in\rho_{0}$ and $\rho^{\ast}(x,T)\in\rho_{T}$ and in the
same time satisfying
$\mathcal{S}_{\Sigma}(\rho^{\ast},\theta^{\ast})<\mathcal{S}_{\Sigma}(\rho,\theta)$
for every $\rho\in Q$ and $\theta\in\mathcal{U}$. Here, $\mathcal{S}_{\Sigma}$
is the cost functional of ${\Sigma}$ and it is defined by

$\mathcal{S}_{\Sigma}=\int_{0}^{T}\ell(\rho(x,t),\theta(x,t))dt+\Phi(\rho(x,T),\rho_{T})$
(2.23)

where $\ell:\mathcal{U}\times\mathbb{R}\mapsto\mathbb{R}$ is the cost
Lagrangian and $\Phi:\mathrm{Dens}(Q)\times\mathrm{Dens}(Q)\to\mathbb{R}$ is
the cost for the final state.

To solve the optimal control problem Pontryagin’s Maximum Principle is used
and the cost functional (2.23) is now written in terms of Hamilton’s principle

$\mathcal{S}_{\Sigma}=\int_{t_{0}}^{t_{1}}\left(\left\langle\lambda,\dot{\rho}\right\rangle-
H_{\Sigma}(\lambda,\theta)\right)dt$ (2.24)

According to the Maximum/Minimum Principle, we first need to define a control
Hamiltonian, $H_{\Sigma}:\mathcal{U}\times
T^{\ast}\mathrm{Den}(Q)\mapsto\mathbb{R}$, and in this particular case it is

$H_{\Sigma}(\lambda,\theta)=\left\langle\lambda,\pounds_{f(x,\theta)}{\rho(x,t)}-\frac{\nu^{2}}{2}\pounds_{\xi}{\pounds_{\xi}{\rho(x,t)}}\right\rangle-\ell(\theta,\rho).$
(2.25)

###### Remark 2.8.

Taking things one step further, we substitute the Legendre transformation
(2.25) into the action $\mathcal{S}_{\Sigma}$ to obtain

$\mathcal{S}_{\Sigma}=\int_{t_{0}}^{t_{1}}\left(\ell(\theta,\rho)+\left\langle\lambda,\dot{\rho}+\pounds_{f(x,\theta)}{\rho}-\frac{\nu^{2}}{2}\pounds_{\xi}{\pounds_{\xi}{\rho}}\right\rangle\right)dt$
(2.26)

When the control vector field $\theta$ belongs to Lie algebra and the
underlying the dynamical system is described by a Lie group of volume
preserving diffeomorphisms action on elements of $Q$ and the noise terms are
absent, then this becomes the fluid mechanics formulation of Monge–Kantorovich
optimal transport problem of [10]. Our interest in the Lagrangian version with
Lagrange multipliers is due to the fact that it allows us to derive the
multisymplectic structure of the problem as we shall see in section 3 as it
was introduced in [31].

If the following trajectory $(\rho,\theta)$ solves the optimal control problem
of the definition 2.7, then according to Pontryagin’s Maximum Principle there
exists a dual vector field $\chi:[t_{0},t_{1}]\mapsto T^{\ast}\mathrm{Den}(Q)$
along $\rho$ such that

$\frac{d}{dt}\chi=X_{H^{\ast}_{\Sigma}},$

where $X_{H^{\ast}_{\Sigma}}$ is the Hamiltonian vector field associated with
$H^{\ast}_{\Sigma}$ and

$H^{\ast}_{\Sigma}(\lambda)=\sup\left\\{H_{\Sigma}(\lambda,\theta)\quad|\quad\theta\in\mathcal{U}\right\\}.$
(2.27)

In other words, $\chi$ is the integral curve of the Hamiltonian vector field
$X_{H^{\ast}_{\Sigma}}$. Without getting much into technical details, the
maximum Hamiltonian, $H^{\ast}_{\Sigma}$, is locally integrable smooth and the
covector field $\chi$ is locally absolutely continuous. That said, in local
coordinates $(\rho,\lambda)$, $X_{H^{\ast}_{\Sigma}}$ are governed by the
following set of equations

$\displaystyle\partial_{t}{\rho}\,d^{n}x$
$\displaystyle=\frac{\delta}{\delta\lambda}H^{\ast}_{\Sigma}=-\pounds_{f(x,\theta)}{(\rho\,d^{n}x)}+\frac{\nu^{2}}{2}\pounds_{\xi}{\pounds_{\xi}{(\rho\,d^{n}x)}},$
(2.28) $\displaystyle\partial_{t}{\lambda}$
$\displaystyle=-\frac{\delta}{\delta\rho}H^{\ast}_{\Sigma}=\pounds_{f(x,\theta)}{\lambda}-\frac{\nu^{2}}{2}\pounds_{\xi}{\pounds_{\xi}{\lambda}}+\frac{\delta\ell(\theta,\rho)}{\delta\rho},$
$\displaystyle\rho(x,0)$ $\displaystyle=\rho_{0}(x),$
$\displaystyle\lambda(x,T)$
$\displaystyle=\frac{\delta\mathcal{S}_{\Sigma}}{\delta\rho_{T}(x)},$

where $\frac{\delta}{\delta\lambda}$ and $\frac{\delta}{\delta\rho}$ denote
the functional derivative with respect to $\lambda$ and $\rho$ respectively.
Such derivative is defined by

$\frac{\delta H}{\delta\alpha}=\lim\limits_{\epsilon\to
0}\frac{H^{\ast}_{\Sigma}(\alpha+\epsilon\delta\alpha)-H^{\ast}_{\Sigma}(\alpha)}{\epsilon}.$

This system of equations is known as _mean–field games_ [73, 74, 58]. As
mentioned before, the first equation is the Fokker–Planck equation, and the
adjoint equation is the stochastic Hamilton–Jacobi–Bellman equation [113]. The
Lie derivative in this case is for a scalar quantity and in general
coordinates it is

$\displaystyle\pounds_{f(x,\theta)}\lambda=f(x,\theta)\cdot\nabla\lambda.$

The stochastic Hamilton–Jacobi–Bellman equation is a nonlinear backward
equation and it appears as the continuous limit of dynamic programming
solution of stochastic optimal control problem with cost functional

$\displaystyle\mathcal{S}$
$\displaystyle=\int_{0}^{t}\ell(\theta)\,dt+\Phi(X_{t},x_{T}),$

subject to

$\displaystyle
d_{t}X_{t}=f(X_{t},\theta)\,dt+\nu\sum\limits_{i}\xi_{i}(X_{t})dW^{i}_{t},\quad$
$\displaystyle\,X_{0}=x_{0},\quad\text{and}\quad X_{T}=x_{T}.$

The variational principle form of the mean–field games was studied in a number
of publications, mainly [113, 90, 11]. Further, such minimisation is connected
to the Schrödinger bridge problem, where it seeks to find the flow of
probability density that connects an initial density and a final density
distribution while minimising the Kullback-Leibler divergence. In recent
years, Schrödinger bridge was regarded as a way to regularise the optimal
transport problem [91, 80, 79, 29, 28, 103].

The system of coupled equations is difficult to solve analytically, and that
leaves us with a numerical approach. Moreover, an additional challenge arises
because of the backward and nonlinear nature of the stochastic
Hamilton–-Jacobi-–Bellman equation. Here, the adjoint equation is not
discussed in detail, as later in 3 we demonstrate how it can be eliminated
from the system of equations. One numerical solution method was introduced in
[1]. In section 4, we introduce a structure-preserving numerical scheme for
mean–field games. The core element of the scheme is that it relies on
variational principle for its derivation.

##  3 Multisymplectic geometry for deep learning

Extending the philosophy of [122, 82], uncertainties in the form of Wiener
processes are introduced to the dynamical model of deep neural networks and
rendering the training problem as a stochastic optimal control problem. It is
worth pointing out that there are numerous neural network architectures that
can be used for deep learning, but the focus will solely be on residual neural
networks. As in [121], the residual neural networks are formulated as a
dynamical system, but to consider a more general and robust case, the residual
network dynamics are treated as a stochastic system

$\displaystyle
d_{t}X_{t}=f(X_{t},\theta)\,dt+\nu\sum\limits_{k}\xi_{k}(X_{t})dW_{t}^{k},\quad
X_{0}=x_{0},$ (3.1)

where the initial condition $X_{0}\in Q$ is the input to the network and the
final condition, $X_{T}$, is the output. Here, we choose the parameters
$\theta=u(x,t)\in\mathfrak{X}(Q)$, i.e. the network’s weight and it is a
vector field that belong to $\mathfrak{X}(Q)$, which the space of vector
fields on $Q$. This means $\mathfrak{X}(Q)$ is the space of admissible
controls $\mathcal{U}$. The flow generated by vector fields on
$\mathfrak{X}(Q)$ are elements of the group of diffeomorphisms on $Q$,
$\mathrm{Diff}(Q)$. The function $f:Q\times\mathfrak{X}(Q)\to TQ$ is the
activation function.

###### Remark 3.1.

In the language of optimal control, here the network’s weights and biases,
$\theta\in\mathfrak{X}(Q)$, play the role of the control, often denoted in
control theory literature by $u$. This is used in order to distinguish the
network’s parameters, which are coefficients, from the usual optimal control
$u$, which is a vector field.

The problem of training, supervised one, is concerned with finding a suitable
$\theta$ such that output $X_{T}=x_{T}$ would be as close as possible to the
chosen target datum $c\in\mathcal{C}$, where $\mathcal{C}$ is the space of
target data, with respect to a certain metric. Thus, as introduced in [82],
the supervised training can be formulated as an optimal control problem, with
the dataset $\\{(x_{0}^{(i)},c^{(i)})\\}_{i=0,\dots,N}$ specifies the initial
conditions and the final conditions. That means one needs to find $\theta$
that would minimise $\sum\limits_{i=0}^{N}\Phi(\pi(X^{(i)}_{T}),c^{(i)})$,
where $\Phi:\mathcal{C}\times\mathcal{C}\to\mathbb{R}$ is a loss function and
$\pi:Q\to\mathcal{C}$ acts as projection from the state manifold $Q$ to the
space of target data. The supervised learning problem can then be stated as:

###### Problem 3.2.

Let $(Q,\mathcal{F},\mathbb{P})$ be a probability space and the stochastic
control system $\Sigma=(Q,f,\mathfrak{X}(Q))$. Let $x_{0}^{(i)}\in Q$ and
$c^{(i)}\in\mathcal{C}$ for $i=0,\dots,N$. The cost functional is chosen as

$\displaystyle\mathcal{S}_{\Sigma}=\sum_{i=0}^{N}\Phi(\pi(X^{(i)}_{T}),c^{(i)})+\int_{0}^{t}\int_{Q}\widehat{\ell}(\theta)\,d\rho\,dt,$
(3.2)

subject to the following constraint

$\displaystyle
d_{t}X_{t}^{(i)}=f(X_{t}^{(i)},\theta)\,dt+\nu\sum_{k}\xi_{k}(X_{t}^{(i)})dW_{t}^{k},\quad
X^{(i)}_{0}=x^{(i)}_{0},$ (3.3)

where $\widehat{\ell}:\mathfrak{X}(Q)\to\mathbb{R}$ is the cost Lagrangian and
$\xi_{k}\in TQ$ are vector fields used to couple the independent and
identically distributed Brownian motions $W^{k}_{t}$ to the tangent bundle
$TQ$. The training of a deep residual network is done by finding the
controlled trajectories $(x^{(i),\ast},\theta^{\ast})$ with
$X^{(i),\ast}_{0}=x^{(i)}_{0}\in Q$ and $\theta\in\mathfrak{X}(Q)$ such that

$\mathcal{S}_{\Sigma}(\theta^{\ast})<\mathcal{S}_{\Sigma}(\theta),\quad\forall(\theta)\in\mathfrak{X}(Q).$

This stochastic optimal control problem can be formulated in terms of
mean–field games as in [122]. Although, as mentioned earlier, in [122], the
McKean–Vlasov approach to mean–field games was used. Here, we use the
Fokker–Planck and Hamilton–Jacobi–Bellman approach, which seeks to control the
probability density of the stochastic dynamical system directly. With that,
problem 3.2 is equivalent to

###### Problem 3.3.

Let $(Q,\mathcal{F},\mathbb{P})$ be a probability space and the stochastic
control system $\Sigma_{\rho}=(\mathrm{Dens}(Q),f,\mathfrak{X}(Q))$. Let
$\rho^{i}_{0}\in\mathrm{Dens}(Q)$ and $c^{(i)}\in\mathcal{C}$ for
$i=0,\dots,N$. The cost functional is chosen as

$\displaystyle\mathcal{S}_{\Sigma_{\rho}}=\sum\limits_{i=0}^{N}\int_{Q}\Phi(\pi(\rho(x,T)),c^{(i)})d\rho+\int_{0}^{t}\ell(\rho,\theta)\,dt,\quad\text{with}\quad\ell(\rho,\theta)=\int_{Q}\widehat{\ell}(\theta)\,d\rho,$
(3.4)

subject to the following constraint

$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f(x,\theta)}\rho\,d^{n}x-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}\rho\,d^{n}x,\quad\rho(x,0)=\rho_{0}(x),$
(3.5)

where $\widehat{\ell}:\mathfrak{X}(Q)\to\mathbb{R}$ is the cost Lagrangian.
Then, the training of a deep residual network is done by finding the
controlled trajectories $(\rho^{\ast},\theta^{\ast})$ with
$\rho^{\ast}(x,0)=\rho_{0}\in\mathrm{Dens}(Q)$ and $\theta\in\mathfrak{X}(Q)$
such that

$\mathcal{S}_{\Sigma_{\rho}}(\rho^{\ast},\theta^{\ast})<\mathcal{S}_{\Sigma_{\rho}}(\rho,\theta),\quad\forall(\rho,\theta)\in\mathrm{Dens}(Q)\times\mathfrak{X}(Q).$

To obtain the weights $\theta$ that would yield the trajectory
$(\rho^{\ast},\theta^{\ast})$, we apply Pontryagin’s Maximum Principle, with
the Hamiltonian defined by

$\displaystyle H_{\Sigma_{\rho}}\quad:\quad
T^{\ast}\mathrm{Dens}(Q)\times\mathfrak{X}(Q)\mapsto\mathbb{R},$
$\displaystyle
H_{\Sigma_{\rho}}(\rho,\lambda,v)=\ell(\rho,\theta)+\left\langle\lambda,\pounds_{f(x,\theta)}\rho-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}\rho\right\rangle.$
(3.6)

Here $\langle\cdot,\cdot\rangle$ is pairing and it is a mapping
$\langle\cdot,\cdot\rangle:T^{\ast}\mathrm{Dens}(Q)\times
T\mathrm{Dens}(Q)\to\mathbb{R}$. However, we apply Legendre transformation
[84] to switch to the Lagrangian picture as it is needed when constructing
numerical schemes based on variational principle. The Lagrangian is then

$\displaystyle\ell(\rho,\theta)+\left\langle\lambda,\partial_{t}\rho-\pounds_{f(x,\theta)}\rho-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}\rho\right\rangle_{T^{\ast}\mathrm{Dens}(Q)\times
T\mathrm{Dens}(Q)},$ (3.7)

and it is known as Clebsch Lagrangian [53]. The stationary variations of the
action (3.4) with constraints

$\displaystyle\delta\mathcal{S}_{\Sigma_{\rho}}=\sum_{i=0}^{N}\delta\int_{Q}\Phi(\pi(\rho(x,T)),c^{(i)}))d\rho+\delta\int_{0}^{t}\ell(\rho,\theta)\,dt,$
$\displaystyle+\delta\int_{0}^{t}\left\langle\lambda,\partial_{t}\rho-\pounds_{f(x,\theta)}\rho-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}\rho\right\rangle\,dt,$

as in (2.10) yields the following Euler-Lagrange equations:

$\displaystyle\delta\lambda$
$\displaystyle:\quad\partial_{t}\rho\,d^{n}x+\pounds_{f(x,\theta)}(\rho\,d^{n}x)-\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}(\rho\,d^{n}x)=0,$
(3.8) $\displaystyle\delta\rho$
$\displaystyle:\quad\partial_{t}\lambda+\pounds_{f(x,\theta)}\lambda+\frac{\nu^{2}}{2}\sum_{i,j}\pounds_{\xi_{i}}\pounds_{\xi_{j}}\lambda+\frac{\delta\ell}{\delta\rho}=0,$
$\displaystyle\delta\theta$
$\displaystyle:\quad\frac{\delta\ell}{\delta\theta}+\rho\nabla\lambda\,dx\otimes
d^{n}x=0,$

with $\quad\rho(x,0)=\rho_{0}$ and
$\lambda(x,T)=\frac{\delta\Phi}{\delta\rho_{T}}$. The last equation is the
optimality condition. The variational derivative
$\frac{\delta\ell}{\delta\theta}$ is in fact a one-form density with a
differential form $dx\otimes d^{n}x$. Substituting the optimality condition
into the Hamiltonian $H_{\Sigma_{\rho}}$ (3.6), gives us the following
Hamiltonian $H_{\Sigma_{\rho}}^{\ast}:T^{\ast}\mathrm{Dens}(Q)\to\mathbb{R}$
which is now a function in terms of canonical coordinates of the cotangent
bundle $T^{\ast}\mathrm{Dens}(Q)$ and it recovers the system of equations
(2.28).

###### Remark 3.4.

The nomenclature for the partial differential equation for the evolution
Lagrange multiplier $\lambda$ might cause confusion. Despite it being called
the stochastic Hamilton–Jacobi–Bellman, in fact the equation is deterministic.
The reason for the naming is that it is the continuous limit of dynamic
programming when applied to stochastic optimal control problems.

###### Remark 3.5.

The mean–field games system (3.8) becomes fluid mechanics formulation of the
the Monge–Kantorovich optimal transport problem when $\xi_{k}\to 0$ [10]. What
is interesting regarding this formulation is that it relates the solution of
the optimal transport problem to the solution of Euler’s incompressible fluid
equation.

Solving (3.8) yields the training weights $\theta$. Still, finding the
solution is challenging, even for numerical methods and for this reason we
cast it as a multisymplectic system with the aim of using multisymplectic
integrators. Here we are following the same approach to multisymplectic
systems as [31], where the stationary variations of the Clebsch variational
principle for back-to-label inverse mapping yield a multisymplectic system.

Before applying Pontryagin’s Maximum Principle, first, a change of variable is
required to carry out the computation and it is done by decomposing the
weights $\theta:Q\to TQ$ as a sum of two vector fields $w:Q\to TQ$ and osmotic
velocity term, $\frac{\nu^{2}}{2}\nabla\mathrm{log}(\rho(x,t))$ , i.e.

$\displaystyle\theta(x,t)=w(x,t)+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho(x,t)).$

This decomposition was used in [97] to show the relation between stochastic
diffusion process and the Schrödinger equation. It was further used in [29,
28, 103] for Schrödinger bridge problem.

###### Remark 3.6.

The term $\nabla\mathrm{log}(\rho(x,t))=\frac{\nabla\rho(x,t)}{\rho(x,t)}$
contains the density function in the denominator. In this article, the main
assumption is that the probability density consists of a single or a sum of
Gaussian functions. Depending on the configuration manifold $Q$, the
$\rho(x,t)>0$ for all $x\in Q$.

Here, the vector field that describes the residual neural network becomes

$\displaystyle f(x,\theta)=f\left(w(x,t)+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho(x,t))\right),$ (3.9)

where $f$ is the activation function and in this paper, $f$ is a linear
function. Substituting it into the Fokker–Planck equation we get

$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f(w)}(\rho\,d^{n}x)=0,\quad\rho(x,0)=\rho_{0},$
(3.10)

Then the optimal control problem becomes slightly different, as required.

###### Problem 3.7.

Let $(Q,\mathcal{F},\mathbb{P})$ be a probability space and the stochastic
control system $\Sigma_{\rho}=(\mathrm{Dens}(Q),\nabla,\mathfrak{X}(Q))$. Let
$\rho_{0}\in\mathrm{Dens}(Q)$ and $c^{(i)}\in\mathcal{C}$ for $i=0,\dots,N$.
The cost functional is chosen as

$\displaystyle\mathcal{S}_{\Sigma_{\rho}}=\sum_{i=0}^{N}\int_{Q}\Phi(\pi(\rho(x,T)),c^{(i)})d\rho+\int_{0}^{t}\ell\left(\rho,w+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho)\right)\,dt,$ (3.11)

subject to the following constraint

$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f(w)}(\rho\,d^{n}x)=0,\quad\rho(x,0)=\rho_{0}(x),$
(3.12)

where $\ell:T\mathrm{Dens}(Q)\times\mathfrak{X}(Q)\to\mathbb{R}$ is the cost
Lagrangian defined by

$\displaystyle\ell\left(\rho,w+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho)\right)=\int_{Q}\left\|w+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho)\right\|^{2}\,d\rho,$
(3.13)

and $\ell:\mathfrak{X}(Q)\to\mathbb{R}$ is a convex function. Finding the
controlled trajectories $(\rho^{\ast},\theta^{\ast})$ with
$\rho^{\ast}(x,0)=\rho_{0}\in Q$ and $\theta\in\mathfrak{X}(Q)$ such that

$\mathcal{S}_{\Sigma_{\rho}}(\rho^{\ast},\theta^{\ast})<\mathcal{S}_{\Sigma_{\rho}}(\rho,\theta),\quad\forall(\rho,\theta)\in\mathrm{Dens}(Q)\times\mathfrak{X}(Q),$

solves the training of a deep residual network.

The evolution equation of the probability density is now the continuity
equation, which is a first order partial differential equation, and this is
essential to our approach as it allows us to reduce (3.8) to Euler–Poincaré
equation. The result now is that the Hamiltonian is affine in the derivatives
with respect to independent variables, which allows us to obtain a
multisymplectic formulation. With that the Pontryagin’s Hamiltonian becomes

$\displaystyle H_{\Sigma_{\rho}}\quad:\quad
T^{\ast}\mathrm{Dens}(Q)\times\mathfrak{X}(Q)\mapsto\mathbb{R},$
$\displaystyle
H_{\Sigma_{\rho}}(\rho,\lambda,v)=\ell\left(\rho,w+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho)\right)+\left\langle\lambda,\pounds_{f(w)}\rho\right\rangle_{T^{\ast}\mathrm{Dens}(Q)\times
T\mathrm{Dens}(Q)},$

and applying the Legndre transformation, we map the Hamiltonian to a
Lagrangian function, allowing us to use Clebsch variational principle

$\displaystyle 0=\delta\mathcal{S}_{\Sigma_{\rho}}$
$\displaystyle=\delta\int_{0}^{t}\ell\left(\rho,w+\frac{\nu^{2}}{2}\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}(\rho)\right)+\left\langle\lambda,\partial_{t}\rho+\pounds_{f(w)}\rho\right\rangle\,dt$
(3.14)
$\displaystyle+\sum_{i=0}^{N}\delta\int_{Q}\Phi(\pi(\rho(x,T)),c^{(i)})d\rho,$

As in (2.11), the stationary variations yield the Euler-Lagrange system of
partial differential equations

$\displaystyle\delta\lambda$
$\displaystyle:\quad\partial_{t}\rho\,d^{n}x+\pounds_{f(w)}(\rho\,d^{n}x)=0,$
(3.15) $\displaystyle\delta\rho$
$\displaystyle:\quad\partial_{t}\lambda+\pounds_{f(w)}\lambda+\frac{\delta\ell}{\delta\rho}=0,$
$\displaystyle\delta w$ $\displaystyle:\quad\frac{\delta\ell}{\delta
w}+\rho\nabla\lambda\,dx\otimes\,d^{n}x=0,$

with $\rho(x,0)=\rho_{0}$ and $\lambda(x,T)=\frac{\delta\Phi}{\delta\rho_{T}}$
and such system of equations has a multisymplectic structure:

$\displaystyle\begin{bmatrix}0&0&0\\\ 0&0&1\\\
0&-1&0\end{bmatrix}\begin{bmatrix}\partial_{t}w^{i}\\\ \partial_{t}\lambda\\\
\partial_{t}\rho\end{bmatrix}+\begin{bmatrix}0&-\rho\frac{\partial
f}{\partial\zeta}&\frac{\nu^{2}}{2}\\\ \rho\frac{\partial
f}{\partial\zeta}&0&f(w)\\\
-\frac{\nu^{2}}{2}&-f(w)&0\end{bmatrix}\begin{bmatrix}\partial_{x^{i}}w^{i}\\\
\partial_{x^{i}}\lambda\\\ \partial_{x^{i}}\rho\end{bmatrix}=\nabla_{z}S(z),$
(3.16)

where $z=(w^{i},\lambda,\rho)^{T}$ and the Hamiltonian is given by

$\displaystyle
S(z)=\frac{\rho}{2}\left\|w\right\|^{2}+\frac{\nu^{4}\rho}{8}\left\|\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}{\rho}\right\|^{2}.$ (3.17)

The last term in $S(z)$ is the _Fisher information metric_. The structure in
(3.16) is the sum of pre-symplectic structures, each pre-symplectic structure
corresponds to an independent variable and they are all skew-symmetric
matrices. In this case, the pre-symplectic structure corresponding to the
independent variable $t$ and the pre-symplectic structure corresponding to the
spatial independent variable $x^{i}$ are

$\displaystyle M=\begin{bmatrix}0&0&0\\\ 0&0&1\\\
0&-1&0\end{bmatrix},\quad\text{and}\quad
K_{i}=\begin{bmatrix}0&-\rho\frac{\partial
f}{\partial\zeta}&\frac{\nu^{2}}{2}\\\ \rho\frac{\partial
f}{\partial\zeta}&0&f(w)\\\ -\frac{\nu^{2}}{2}&-f(w)&0\end{bmatrix},$ (3.18)

respectively. Writing the system of partial differential equations (3.15) as
(3.16) using anti-symmetric matrices

$\displaystyle
M\partial_{t}z+\sum\limits_{i}K_{i}\partial_{x^{i}}z=\partial_{z}S(z),$ (3.19)

is an alternative approach to multisymplectic partial differential equations
[19]. In fact, according to [19], a system is multisymplectic if it can be
written in the form of (3.19). In a geometric setting, the equations (3.16) of
can be written intrinsically, like in equation (2.11), as

$\displaystyle(j^{1}\phi)^{\ast}\left[j^{1}(V)\mathbin{\hbox to6.0pt{\vrule
height=0.4pt,width=5.0pt,depth=0.0pt\kern-0.4pt\vrule
height=6.0pt,width=0.4pt,depth=0.0pt\hss}}\Omega_{\mathcal{L}}\right]=0,$

where $V\in J^{1}(E)$ and the Lagrangian density
$\mathcal{L}:J^{1}(E)\to\Lambda^{n+1}(Q)$, here is defined by

$\displaystyle\mathcal{L}:=\ell_{c}\,d^{n}x\wedge dt$
$\displaystyle=\frac{\rho}{2}\left\|w\right\|^{2}+\frac{\nu^{2}}{2}\rho\left\langle
w,\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}{\rho}\right\rangle+\frac{\nu^{4}\rho}{8}\left\|\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}{\rho}\right\|^{2}$
$\displaystyle+\left\langle\lambda,\partial_{t}\rho+\pounds_{f(w)}\rho\right\rangle\,d^{n}x\wedge
dt.$

With the Cartan form $\Theta_{\mathcal{L}}\in\Lambda^{n+1}(J^{1}(E))$, and its
exterior derivative
$\Omega_{\mathcal{L}}=-d\Theta_{\mathcal{L}}\in\Lambda^{n+2}(J^{1}(E))$, is

$\displaystyle\Theta_{\mathcal{L}}$
$\displaystyle=\frac{\partial\ell_{c}}{\partial\rho_{x_{i}}}d\rho\wedge\left(\partial_{x_{i}}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n}x\right)\wedge
dt+\frac{\partial\ell_{c}}{\partial\rho_{t}}d\rho\wedge
d^{n}x+\frac{\partial\ell_{c}}{\partial
w^{j}_{x_{i}}}dw^{j}\wedge\left(\partial_{x_{i}}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}d^{n}x\right)\wedge
dt.$ (3.20)

Without the loss of generality, when the base manifold $Q$ is
$\mathbb{R}^{2}\times\mathbb{R}$ with coordinates $(x,y,t)$ and the fibres are
$w=(w^{1},w^{2})$, the form $\Theta_{\mathcal{L}}$ is expressed explicitly by

$\displaystyle\Theta_{\mathcal{L}}$
$\displaystyle=\left(\frac{\nu^{2}}{2}w^{1}+\frac{\nu^{4}}{4}\frac{\partial_{x}\rho}{\rho}+\lambda
w^{1}\right)d\rho\wedge dx^{2}\wedge
dt+\left(\frac{\nu^{2}}{2}w^{2}+\frac{\nu^{4}}{4}\frac{\partial_{y}\rho}{\rho}+\lambda
w^{2}\right)d\rho\wedge dx^{1}\wedge dt$ $\displaystyle+\lambda d\rho\wedge
dx^{1}\wedge dx^{2}+\left(\lambda\rho\right)dw^{1}\wedge dx^{2}\wedge
dt+\left(\lambda\rho\right)dw^{2}\wedge dx^{1}\wedge
dt+\left(\ell-\frac{\partial\ell_{c}}{\partial\rho_{x}}\partial_{x}\rho\right.$
$\displaystyle\left.-\frac{\partial\ell_{c}}{\partial\rho_{y}}\partial_{y}\rho-\frac{\partial\ell_{c}}{\partial\rho_{t}}\partial_{t}\rho-\frac{\partial\ell_{c}}{\partial
w^{1}_{x}}\partial_{x}w^{1}-\frac{\partial\ell_{c}}{\partial
w^{2}_{x}}\partial_{x}w^{2}\right)dx^{1}\wedge dx^{2}\wedge dt.$

The Hamilton–Jacobi–Bellman equation, whose solution is $\lambda$, in (3.16)
can be eliminated and the result is an equivalent system that is governed by
the Euler–Poincaré equation [56]. For multisymplectic equations, the
elimination was presented in [31] and here, we follow the same steps to obtain
the Euler–Poincaré equation for deep learning.

###### Theorem 3.8.

Consider the optimal control problem 3.7 for supervised learning of stochastic
residual network, it’s solution, denoted by the tuple
$\left(w,\lambda,\rho\right)$, satisfies

$\displaystyle(j^{1}\phi)^{\ast}\left[j^{1}(V)\mathbin{\hbox to6.0pt{\vrule
height=0.4pt,width=5.0pt,depth=0.0pt\kern-0.4pt\vrule
height=6.0pt,width=0.4pt,depth=0.0pt\hss}}\Omega_{\mathcal{L}}\right]=0,$

where $\Omega_{\mathcal{L}}=-d\Theta_{\mathcal{L}}$ of (3.20). Then,
$\frac{\delta\ell}{\delta w}=\rho
w+\frac{\nu^{2}}{2}\nabla\rho=\rho\nabla\lambda$ and $\rho$ satisfy the
following Euler–Poincaré equation on $\mathfrak{X}(Q)\times\mathrm{Dens}(Q)$

$\displaystyle\left[\frac{\partial}{\partial t}\frac{\delta\ell}{\delta
w}+\mathrm{ad}^{\ast}_{f(w)}\frac{\delta\ell}{\delta
w}-\rho\diamond\frac{\delta\ell}{\delta\rho}\right]_{i}dx^{i}\otimes
d^{n}x=0,$ (3.21)
$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f(w)}(\rho\,d^{n}x)=0.$

###### Proof.

This elimination is done by taking the time derivative of the inner product of
$\frac{\delta\ell}{\delta w}$ and an arbitrary element
$\eta\in\mathfrak{X}(Q)$. The inner product is defined by

$\displaystyle\left\langle\frac{\delta\ell}{\delta
w},\eta\right\rangle_{\mathfrak{X}^{\ast}(Q)\times\mathfrak{X}(Q)}=\int_{Q}\frac{\delta\ell}{\delta
w}\cdot\eta\,\,d^{n}x.$ (3.22)

With the inner product defined and using the equations for $\partial_{t}\rho$
and $\partial_{t}\lambda$, the time derivative of $\frac{\delta\ell}{\delta
w}$ is

$\displaystyle\frac{d}{dt}\left\langle\frac{\delta\ell}{\delta
w},\eta\right\rangle$
$\displaystyle=\frac{d}{dt}\left\langle\rho\nabla\lambda,\eta\right\rangle$
$\displaystyle=-\left\langle\pounds_{f(w)}\rho,(\eta\cdot\nabla)\lambda\right\rangle+\left\langle\rho,(\eta\cdot\nabla)\pounds_{f(w)}\lambda+(\eta\cdot\nabla)\frac{\delta\ell}{\delta\rho}\right\rangle$
$\displaystyle=-\left\langle\frac{\delta\ell}{\delta
w},\mathrm{ad}_{f(w)}\eta\right\rangle+\left\langle\rho\diamond\frac{\delta\ell}{\delta\rho},\eta\right\rangle$
$\displaystyle=-\left\langle\mathrm{ad}^{\ast}_{f(w)}\frac{\delta\ell}{\delta
w},\eta\right\rangle+\left\langle\rho\diamond\frac{\delta\ell}{\delta\rho},\eta\right\rangle.$

Since $\eta$ is arbitrary, in the end we obtain the Euler–Poincaré equation

$\displaystyle\left[\frac{\partial}{\partial t}\frac{\delta\ell}{\delta
w}+\mathrm{ad}^{\ast}_{f(w)}\frac{\delta\ell}{\delta
w}-\rho\diamond\frac{\delta\ell}{\delta\rho}\right]_{i}\,dx^{i}\otimes
d^{n}x=0,$ (3.23)
$\displaystyle\partial_{t}\rho\,d^{n}x+\pounds_{f(w)}(\rho\,d^{n}x)=0.$

∎

The diamond operator, introduced in [56],
$\diamond:T^{\ast}\mathrm{Dens}(Q)\to\mathfrak{X}^{\ast}(Q)$ is defined, for
any $\eta\in\mathfrak{X}(Q)$, by

$\displaystyle\left\langle\rho\diamond\lambda,\eta\right\rangle_{\mathfrak{X}^{\ast}(Q)\times\mathfrak{X}(Q)}=-\left\langle\lambda,\pounds_{\eta}\rho\right\rangle_{T^{\ast}\mathrm{Dens}(Q)\times
T\mathrm{Dens}(Q)}.$ (3.24)

In terms of (3.23), it is
$\rho\diamond\frac{\delta\ell}{\delta\rho}=\rho\nabla\frac{\delta\ell}{\delta\rho}$.
Here the map
$\mathrm{ad}^{\ast}:\mathfrak{X}(Q)\times\mathfrak{X}^{\ast}(Q)\to\mathfrak{X}^{\ast}(Q)$
is the coadjoint map on $\mathfrak{X}^{\ast}(Q)$, and it is obtained by taking
the dual of the adjoint map
$\mathrm{ad}:\mathfrak{X}(Q)\times\mathfrak{X}(Q)\to\mathfrak{X}(Q)$. The
adjoint map is defined as the commutator of two elements, $w,v$, of the Lie
algebra $\mathfrak{X}(Q)$

$\displaystyle[w,v]=\mathrm{ad}_{w}v=\left(w\cdot\nabla\right)v-\left(v\cdot\nabla\right)w,$

where the bracket
$[\cdot,\cdot]:\mathfrak{X}(Q)\times\mathfrak{X}(Q)\to\mathfrak{X}(Q)$ is the
Lie bracket of the Lie algebra $\mathfrak{X}(Q)$. To move to the dual of Lie
algebra is done using the pairing between $\mathfrak{X}^{\ast}(Q)$ and
$\mathfrak{X}(Q)$ which is actualised by the inner product map
$\langle\cdot,\cdot\rangle:\mathfrak{X}^{\ast}(Q)\times\mathfrak{X}(Q)\to\mathbb{R}$
defined in (3.22). This allows us to identify the coadjoint map

$\displaystyle\left\langle\mathrm{ad}^{\ast}_{w}m,v\right\rangle=\left\langle
m,\mathrm{ad}_{w}v\right\rangle,\quad\forall m\in\mathfrak{X}^{\ast}(Q),$

and using integration by parts of the integral (3.22) that gives an expression
for the coadjoint map

$\displaystyle\mathrm{ad}^{\ast}_{w}m=\left(w\cdot\nabla\right)m+(\nabla
w)^{T}\cdot m+m\left(\nabla\cdot w\right),$ (3.25)

and in coordinates in $\mathbb{R}^{N}$,

$\displaystyle\left(\mathrm{ad}^{\ast}_{w}m\right)_{i}=\frac{\partial}{\partial
x^{j}}(m_{i}w^{j})+m_{j}\frac{\partial w^{j}}{\partial x^{i}}.$ (3.26)

With that, the equation (3.23) is expressed in coordinates on $\mathbb{R}^{N}$
as

$\displaystyle\partial_{t}$ $\displaystyle\left(\frac{\delta\ell}{\delta
w^{i}}\right)+\frac{\partial}{\partial
x^{j}}\left(\left(\frac{\delta\ell}{\delta
w^{i}}\right)f(w)^{j}\right)+\left(\frac{\delta\ell}{\delta
w^{j}}\right)\frac{\partial f(w)^{j}}{\partial
x^{i}}=\rho\frac{\partial}{\partial x^{i}}\frac{\delta\ell}{\delta\rho},$
(3.27) $\displaystyle\partial_{t}$ $\displaystyle\rho+\frac{\partial}{\partial
x^{j}}\left(\rho f(w)^{j}\right)=0.$

This allows us to eliminate the stochastic Hamilton–Jacobi–Bellman equation
which is a backward partial differential equation in (3.15). Furthermore,
both, the Euler–Poincaré equation with and the continuity equation (3.21) are
forward in time differential equations and they describe dynamics on the dual
of the semidirect product Lie algebra
$\mathfrak{s}^{\ast}=\mathfrak{X}^{\ast}(Q)\times\mathrm{Dens}(Q)$. The system
of equations (3.21) does appear when applying Euler–Poincaré variational
principle

$\displaystyle\delta\int_{0}^{t}\int_{Q}\frac{1}{2}\rho w\cdot
w+\frac{\nu^{2}}{2}w\cdot\nabla\rho+\frac{\nu^{4}}{8}\frac{\nabla\rho\cdot\nabla\rho}{\rho}\,dx\,dt=0,$

and the variations of $u$ and $\rho$ has the form

$\displaystyle\delta u=\dot{\eta}+[u,\eta],$
$\displaystyle\quad\delta\rho=\nabla\cdot(\rho\eta)d^{n}x,$

where $\eta\in\mathfrak{X}(Q)$ is an arbitrary vector field with fixed
endpoints at the boundary set to zero. The system of equations (3.27) can be
solved using geometric numerical methods, such as multisymplectic method, and
that is the goal of section 4.

One important property that emerges from the necessary condition for
optimality is in fact a momentum map for the action of the group of
diffeomorphisms on the probability density function associated with stochastic
ResNets. This agrees with the Clebsch [85] framework for dynamics and optimal
control [41].

###### Theorem 3.9.

The stationary variation of the control vector field $w$ defines an
equivariant momentum map

$\displaystyle\mathbf{J}:T^{\ast}\mathrm{Dens}(Q)\to\mathfrak{X}^{\ast}(Q)$
(3.28) $\displaystyle\mathbf{J}(\rho,\lambda)=\frac{\delta\ell}{\delta
w}=\rho\nabla\lambda$

###### Proof.

To demonstrate that $\mathbf{J}=\rho\nabla\lambda$ is a momentum map, we begin
by defining a function $J_{\beta}:T^{\ast}\mathrm{Dens}(Q)\to\mathbb{R}$,
explicitly

$\displaystyle
J_{\beta}=\left\langle\mathbf{J}(\rho,\lambda),\beta\right\rangle_{\mathfrak{X}^{\ast}\times\mathfrak{X}}=\int\rho(\partial_{x^{i}}\lambda)\beta^{i}\,dx,$
(3.29)

where $\beta\in\mathfrak{X}(Q)$ and
$\langle\cdot,\cdot\rangle:\mathfrak{X}^{\ast}\times\mathfrak{X}\to\mathbb{R}$
is a natural pairing between $\mathfrak{X}^{\ast}$ and $\mathfrak{X}$. For all
$F\in T^{\ast}\mathrm{Dens}(Q)$, we compute the Poisson bracket on
$T^{\ast}\mathrm{Dens}(Q)$ of $F$ and $J_{\beta}$ and it is

$\displaystyle\\{F,J_{\beta}\\}=\int_{Q}-\frac{\delta
F}{\delta\rho}\partial_{x^{j}}(\rho\beta^{j})-\rho\frac{\delta
F}{\delta\lambda}(\partial_{x^{i}}\lambda)\beta^{i}\,dx$ (3.30)

The right–hand side can be stated as

$\displaystyle\\{F,J_{\beta}\\}=\left\langle\mathbf{d}F,X_{J_{\beta}}\right\rangle,$
(3.31)

where $X_{J_{\beta}}$ is a vector field given by

$\displaystyle
X_{J_{\beta}}=\left(-\partial_{x^{i}}(\rho\beta^{i}),-\partial_{x^{i}}\lambda\beta^{i}\right)=\left(\frac{\partial
J_{\beta}}{\partial\lambda},-\frac{\partial J_{\beta}}{\partial\rho}\right),$

And this vector field is corresponds to the infinitesimal generator of the
action of $\mathrm{Diff}(Q)$, whose associated vector field is $\beta$, on
$T^{\ast}\mathrm{Den}(Q)$ and it is a Hamiltonian vector, which verifies that
(3.29) holds and thus $\mathbf{J}$ is a momentum map.

To show that the momentum map is equivariant, i.e. for any
$g\in\mathrm{Diff}(Q)$ group action on all $(\rho,\lambda)$ the following
holds

$\displaystyle\mathbf{J}(g\cdot\rho,g\cdot\lambda)=\mathrm{Ad}^{\ast}_{g^{-1}}\mathbf{J}(\rho,\lambda).$
(3.32)

Pairing $\mathbf{J}(g\cdot\rho,g\cdot\lambda)$ with any
$\beta\in\mathfrak{X}(Q)$, with $g(0)=e$ and
$\dot{g}(0)=\eta\in\mathfrak{X}(Q)$. Taking the time derivative of the
left–hand side, we have

$\displaystyle\left.\frac{d}{dt}\right|_{t=0}J_{\beta}(g\cdot\rho,g\cdot\lambda)$
$\displaystyle=\left.\frac{d}{dt}\right|_{t=0}\int(g^{\ast}\rho)(\partial_{x^{i}}g^{\ast}\lambda)\beta^{i}\,dx$
$\displaystyle=\int\partial_{x^{j}}(\rho\eta^{j})(\partial_{x^{i}}\lambda)\beta^{i}\,dx+\int\rho\partial_{x^{j}}((\partial_{x^{i}}\lambda)\eta^{i})\beta^{j}\,dx$
$\displaystyle=-\int\rho\partial_{x^{j}}((\partial_{x^{i}}\lambda)\beta^{i})\eta^{j}\,dx+\int\rho\partial_{x^{j}}((\partial_{x^{i}}\lambda)\eta^{i})\beta^{j}\,dx$

Expanding the gradient of the dot products and rearranging terms, the time
derivative reduces to

$\displaystyle\left.\frac{d}{dt}\right|_{t=0}J_{\beta}(g\cdot\rho,g\cdot\lambda)$
$\displaystyle=\int\rho(\partial_{x^{j}}\lambda)\left[\beta^{i}\partial_{x^{i}}\eta^{j}-\eta^{i}\partial_{x^{i}}\beta^{j}\right]\,dx$

Thus we have $\\{J_{\eta},J_{\beta}\\}=J_{[\eta,\beta]}(\rho,\lambda)$, which
is known as the infinitesimal equivariance. Also, it amounts to the time
derivative of $J_{\mathrm{Ad}_{g^{-1}}\beta}(\rho,\lambda)$, which we compute
explicitly

$\displaystyle\left.\frac{d}{dt}\right|_{t=0}J_{\mathrm{Ad}_{g^{-1}}\beta}(\rho,\lambda)$
$\displaystyle=\left.\frac{d}{dt}\right|_{t=0}\int(\mathrm{Ad}^{\ast}_{g^{-1}}(\rho(\partial_{x^{i}}\lambda)))\beta^{i}\,dx=-\int\mathrm{ad}^{\ast}_{\eta}(\rho(\partial_{x^{i}}\lambda)))\beta^{i}\,dx$
$\displaystyle=-\int\rho(\partial_{x^{j}}\lambda)(\mathrm{ad}_{\beta}\eta)^{j}$
$\displaystyle=J_{[\eta,\beta]}(\rho,\lambda),$

thus we have
$J_{\beta}(g\cdot\rho,g\cdot\lambda)=\mathrm{Ad}^{\ast}_{g^{-1}}J_{\beta}(\rho,\lambda)$
and that implies

$\mathbf{J}(g\cdot\rho,g\cdot\lambda)=\mathrm{Ad}^{\ast}_{g^{-1}}\mathbf{J}(\rho,\lambda),$

which proves that the momentum map
$\mathbf{J}:T^{\ast}\mathrm{Dens}(Q)\to\mathfrak{X}^{\ast}(Q)$ is an
equivariant momentum map. ∎

##  4 Multisymplectic variational integrators for deep learning

Often case in optimal control problems, one is presented with the choice
between an optimise-first-then-discretise approach or discretise-first-then-
optimise. In fact, the methodology presented here does not deviate from
discrete mechanics for optimal control of [64, 99, 33], where a discrete
variational principle was used to translate the optimal control problem into a
nonlinear programming one. This means we discretise-first-then-optimise

The family of multisymplectic integrators were presented in [86, 20]. There
are two ways to obtain such numerical scheme: discretising the variational
integrator, and discretising the abstract form of multisymplectic partial
differential equation [20, 109]. Multisymplectic variational integrators were
first developed in [86] and many discrete analogues such discrete Cartan one-
forms, and discrete Noether’s theorem were also introduced. Then in [81]
further generalisations regarding discretisation of the base manifold were
investigated, others introduced moving mesh points on the base manifold [118].
Variational integrators were applied to various nonlinear differential
equations, such as the nonlinear Schroödinger equation [27], the Camassa-Holm
equation [70], Korteweg-de Vries [57] and the beam equation in
$\mathbb{R}^{3}$ [35, 36]. On the other hand, multisymplectic integrators can
be constructed by directly discretising the Hamiltonian partial differential
equation as in [20, 109, 108, 20, 93, 8, 60, 30]. A nice property of such
integrators is that there is a heuristic approach to constructing modified
multisymplectic Hamiltonian for backward error analysis, as it was shown in
[92, 93, 59].

Here we only focus on variational integrators based on [86]. Variational
integrators are, namely, based on Veselov discretisation for mechanics [119],
and many properties and theorems, such as Noether’s theorem and conservation
laws, have a discrete equivalent. Without the loss of generalisation, here we
consider problems with two dimensions for space and one dimension for time,
hence the base manifold is three dimensional manifold. The first step in
constructing a variational integrator is to discretise the base manifold $Q$
as $\mathbb{B}:=\mathbb{Z}\times\mathbb{Z}\times\mathbb{Z}=\\{(i,j,n)\\}$ and
then the with that the fibre bundle $\mathbb{E}=\mathbb{B}\times\mathcal{F}$,
where $\mathcal{F}$ is a smooth manifold. The quantity $y_{i,j,n}$ is an
element of $\mathbb{E}$ over the point $(i,j,n)$. Since the base manifold now
is a discrete space, that requires the partial derivatives with respect to
independent variables $(x^{1},x^{2},t)$ and to be replaced with finite
differences. For our application, midpoint approximation in time and space is
used and discretisation renders the grid as a collection of right-angled
pyramids with triangular bases, which a pyramid is defined by the four points

$\Delta^{n}_{i,j}=((i,j,n),(i,j,n+1),(i+1,j,n),(i,j+1,n)).$

The set of all pyramids is defined as

$\displaystyle\mathbb{B}:=\left\\{\Delta_{i,j,n}\quad|\quad(i,j,n)\in(0,\dots,N_{x^{1}})\times(0,\dots,N_{x^{2}})\times(0,\dots,N_{t})\right\\}.$

With that, the first jet bundle becomes

$\displaystyle J^{1}(\mathbb{E})$
$\displaystyle=\left\\{\left.(\phi_{i,j,n},\phi_{i,j,n+1},\phi_{i+1,j,n},\phi_{i,j+1,n})\right|(i,j,n)\in\mathbb{B},\phi_{i,j,n},\phi_{i,j,n+1},\phi_{i+1,j,n},\phi_{i,j+1,n}\in\mathcal{F}\right\\}$
$\displaystyle=\mathbb{B}\times\mathcal{F}^{4},$

we note that $(\phi_{i,j,n},\phi_{i,j,n+1},\phi_{i+1,j,n},\phi_{i,j+1,n})\in
J^{1}(\mathbb{E})$ are basically $j^{1}(\phi(\bar{x}))$, where $\bar{x}_{i,n}$
is the centre of $\Delta_{i,j,n}$ and $\phi$ is a smooth section of $E$. As it
is in discrete mechanics, the starting point of constructing a variational
integrator is the discretisation of the action functional. The discrete action
$\mathcal{S}_{d}$ is given by

$\displaystyle\mathcal{S}_{d}$
$\displaystyle=\sum\limits_{(i,j,n)\in\mathbb{B}}\widehat{L}\left(\Delta_{i,j,n},w^{1}_{i,j,n},w^{2}_{i,j,n},\lambda_{i,j,n},\rho_{i,j,n}\right)\Delta
x^{1}\Delta x^{2}\Delta t,$

where $\widehat{L}:J^{1}(\mathbb{E})\to\mathbb{R}$ is the discrete Lagrangian
and it is defined by

$\displaystyle\widehat{L}$ $\displaystyle=\int_{t}^{t+\Delta
t}\int_{x^{2}}^{x^{2}+\Delta x^{2}}\int_{x^{1}}^{x^{1}+\Delta
x^{1}}\frac{\rho}{2}\left\|w\right\|^{2}+\frac{\nu^{2}}{2}\rho\left\langle
w,\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}{\rho}\right\rangle$
$\displaystyle+\frac{\nu^{4}\rho}{8}\left\|\left(\frac{\partial
f}{\partial\zeta}\right)^{-1}\nabla\mathrm{log}{\rho}\right\|^{2}+\left\langle\lambda,\partial_{t}\rho+\pounds_{f(w)}\rho\right\rangle
dx^{1}\,dx^{2}\,dt.$

Here, and henceforth, we use $\Delta_{i,j,n}$ as an argument as a short-hand
notation for the tuple

$\displaystyle\widehat{L}\left(\Delta_{i,j,n},w^{1}_{i,j,n},w^{2}_{i,j,n},\lambda_{i,j,n},\rho_{i,j,n}\right)$
$\displaystyle=\widehat{L}\left(w^{1}_{i,j,n},w^{1}_{i,j,n+1},w^{1}_{i+1,j,n},w^{1}_{i,j+1,n},\right.$
$\displaystyle\left.w^{2}_{i,j,n},w^{2}_{i,j,n+1},w^{2}_{i+1,j,n},w^{2}_{i,j+1,n},\lambda_{i,j,n},\lambda_{i,j,n+1},\right.$
$\displaystyle\left.\lambda_{i+1,j,n},\lambda_{i,j+1,n},\rho_{i,j,n},\rho_{i,j,n+1},\rho_{i+1,j,n},\rho_{i,j+1,n}\right).$

Further, to evaluate the Lagrangian on the discrete first jet bundle
$J^{1}(\mathbb{E})$, we require the use of discrete jet prolongations. There
are a number of ways to do so, and here we use a centred discretisation and
with that for the pyramid $\Delta_{i,j,n}$ and the section $\phi\in
J^{1}(\mathbb{E})$, the space and time derivatives are replaced by

$\displaystyle\frac{\partial\phi}{\partial x^{1}}$
$\displaystyle\approx\frac{\phi_{i+1,j+1/2,n+1/2}-\phi_{i,j+1/2,n+1/2}}{2\Delta
x^{1}},\quad\frac{\partial\phi}{\partial
x^{2}}\approx\frac{\phi_{i+1/2,j+1,n+1/2}-\phi_{i+1/2,j,n+1/2}}{2\Delta
x^{2}},$ $\displaystyle\frac{\partial\phi}{\partial t}$
$\displaystyle\approx\frac{\phi_{i+1/2,j+1/2,n+1}-\phi_{i+1/2,j+1/2,n}}{2\Delta
t}.$

Given the discrete action

$\displaystyle\mathcal{S}_{d}$
$\displaystyle=\sum\limits_{(i,j,n)\in\mathbb{B}}\widehat{L}\left(\Delta_{i,j,n},w^{1}_{i,j,n},w^{2}_{i,j,n},\lambda_{i,j,n},\rho_{i,j,n}\right)\Delta
x^{1}\Delta x^{2}\Delta t,$

here, the Lagrangian on the discrete $J^{1}(\mathbb{E})$ becomes

$\displaystyle\widehat{L}$
$\displaystyle=\frac{1}{2}\bar{\rho}_{i,j,n}((\bar{w}^{1})_{i,j,n})^{2}+\frac{1}{2}\bar{\rho}_{i,j,n}((\bar{w}^{2})_{i,j,n})^{2}+\frac{\nu^{2}}{2(\Delta
x^{2})}\left(\rho_{i+1/2,j+1,n+1/2}\right)(\bar{w}^{2})_{i,j,n}$ (4.1)
$\displaystyle-\frac{\nu^{2}}{2(\Delta
x^{2})}\left(\rho_{i+1/2,j,n+1/2}\right)(\bar{w}^{2})_{i,j,n}+\frac{\nu^{2}}{2(\Delta
x^{1})}\left(\rho_{i+1,j+1/2,n+1/2}-\rho_{i,j+1/2,n+1/2}\right)(\bar{w}^{1})_{i,j,n}$
$\displaystyle+\frac{\nu^{4}}{8}\left(\frac{\left(\rho_{i+1/2,j+1,n+1/2}-\rho_{i+1/2,j,n+1/2}\right)^{2}}{(\Delta
x^{2})\bar{\rho}_{i,j,n}}+\frac{\left(\rho_{i+1,j+1/2,n+1/2}-\rho_{i,j+1/2,n+1/2}\right)^{2}}{(\Delta
x^{1})\bar{\rho}_{i,j,n}}\right)$
$\displaystyle+\bar{\lambda}_{i,j,n}\left(\frac{\rho_{i+1/2,j+1/2,n+1}-\rho_{i+1/2,j+1/2,n}}{\Delta
t}+\frac{\rho_{i+1,j+1/2,n+1/2}w^{1}_{i+1,j+1/2,n+1/2}}{(\Delta
x^{1})}\right.$
$\displaystyle\left.-\frac{\rho_{i,j+1/2,n+1/2}w^{1}_{i,j+1/2,n+1/2}}{(\Delta
x^{1})}+\frac{\rho_{i+1/2,j+1,n+1/2}w^{2}_{i+1/2,j+1,n+1/2}}{(\Delta
x^{2})}\right.$
$\displaystyle\left.-\frac{\rho_{i+1/2,j,n+1/2}w^{2}_{i+1/2,j,n+1/2}}{(\Delta
x^{2})}\right).$

Here, the bar over the variables
$(\bar{w}^{1})_{i,j,n},(\bar{w}^{2})_{i,j,n},\bar{\rho}_{i,j,n},\bar{\lambda}_{i,j,n}$
are used to denote the average of vertices of $\Delta_{i,j,n}$. Similar to the
previous schemes, we take the variations of the discrete action and set it
equal to zero, i.e. $\delta\mathcal{S}_{d}=0$,

$\displaystyle 0$ $\displaystyle=\delta\mathcal{S}_{d}$
$\displaystyle=\sum\limits_{(i,j,n)\in\mathbb{B}}\sum\limits_{\alpha_{i},\alpha_{j},\alpha_{n}=-1}^{1}\left\langle\frac{\partial\hat{L}}{\partial(w^{1})_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}},\delta
w^{1}_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}\right\rangle$
$\displaystyle+\left\langle\frac{\partial\hat{L}}{\partial(w^{2})_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}},\delta
w^{2}_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}\right\rangle+\left\langle\frac{\partial\hat{L}}{\partial\rho_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}},\delta\rho_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}\right\rangle$
$\displaystyle+\left\langle\delta\lambda_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}},\frac{\partial\hat{L}}{\partial\lambda_{i+\alpha_{i},j+\alpha_{j},k+\alpha_{n}}}\right\rangle.$

Here we use the notation
$D_{\alpha_{i},\alpha_{j},\alpha_{n},z}\hat{L}(\phi_{d}(\Delta_{i,j,n}))$
indicates the derivative of $\hat{L}(\phi_{d}(\Delta_{i,j,n}))$ with respect
to
$\phi_{d}(\Delta_{i,j,n})=(w^{1}_{i,j,n},w^{2}_{i,j,n},\rho_{i,j,n},\lambda_{i,j,n})$.
Shifting the indices, we factor out $\delta w^{1}_{i,j,n}$, $\delta
w^{2}_{i,j,n}$, $\delta\rho_{i,j,n}$, and $\delta\lambda_{i,j,n}$ and for
$\delta\mathcal{S}_{d}=0$ and we obtain the _discrete Euler-Lagrange equation_
[86]:

$\displaystyle\sum\limits^{1}_{\alpha_{i},\alpha_{j},\alpha_{n}=-1}\frac{\partial\hat{L}(\phi_{d}(\Delta_{i-\alpha_{i},j-\alpha_{j},n-\alpha_{n}}))}{\partial\phi_{d}(\Delta_{i+\alpha_{i},j+\alpha_{j},n+\alpha_{n}})}$
$\displaystyle=0,$ (4.2)

for $i=1\dots N_{x}$, $j=1\dots N_{y}$ and $n=1\dots N_{t}$. To illustrate how
the discrete variations are obtained, and without the loss of generality, we
consider the one-dimensional of this problem, where the Lagrangian is given by

$\displaystyle\widehat{L}$
$\displaystyle=\frac{1}{2}\bar{\rho}_{i,n}(\bar{w}_{i,n})^{2}+\frac{\nu^{2}}{2}\bar{w}_{i,n}\left(\frac{\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2}}{2\Delta
x}\right)$
$\displaystyle-\frac{\nu^{4}}{8}\bar{\rho}_{i,n}\left(\frac{\rho_{i+1,n+1/2}-2\rho_{i,n+1/2}+\rho_{i-1,n+1/2}}{\bar{\rho}_{i,n}(\Delta
x)^{2}}-\frac{1}{8}\frac{(\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2})^{2}}{(\bar{\rho}_{i,n})^{2}(\Delta
x)^{2}}\right)$
$\displaystyle+\bar{\lambda}\left(\frac{\rho_{i+1/2,n+1}-\rho_{i+1/2,n-1}}{2\Delta
t}+\frac{\rho_{i+1,n+1/2}w_{i+1,n+1/2}-\rho_{i-1,n+1/2}w_{i-1,n+1/2}}{2\Delta
x}\right).$

For the $0=\delta\mathcal{S}_{d}$ to hold, we require the following equations
to be satisfied:

$\displaystyle 0$
$\displaystyle=\frac{\rho_{i,n-1/2}w_{i,n-1/2}}{2}+\frac{\rho_{i,n+1/2}w_{i,n+1/2}}{2}+\frac{\nu^{2}}{4}\left(\frac{\rho_{i+1,n-1/2}+\rho_{i+1,n+1/2}}{2\Delta
x}\right.$ (4.3)
$\displaystyle\left.-\frac{\rho_{i-1,n-1/2}+\rho_{i-1,n+1/2}}{2\Delta
x}\right)-\frac{1}{4}\left(\frac{\lambda_{i+1,n-1/2}\rho_{i,n-1/2}+\lambda_{i+1,n+1/2}\rho_{i,n+1/2}}{\Delta
x}\right.$
$\displaystyle\left.-\frac{\lambda_{i-1,n-1/2}\rho_{i,n-1/2}+\lambda_{i-1,n+1/2}\rho_{i,n+1/2}}{\Delta
x}\right),$ $\displaystyle 0$
$\displaystyle=\frac{\rho_{i+1/2,n+1}-\rho_{i+1/2,n-1}}{2\Delta
t}+\frac{\rho_{i+1,n-1/2}w_{i+1,n-1/2}-\rho_{i-1,n-1/2}w_{i-1,n-1/2}}{4\Delta
x}$
$\displaystyle+\frac{\rho_{i+1,n+1/2}w_{i+1,n+1/2}-\rho_{i-1,n+1/2}w_{i-1,n+1/2}}{4\Delta
x},$ $\displaystyle 0$
$\displaystyle=-\frac{\lambda_{i,n+1}-\lambda_{i,n-1}}{2\Delta
t}-\frac{\nu^{2}}{4}\frac{w_{i+1,k-1/2}+w_{i+1,k+1/2}-w_{i-1,k-1/2}-w_{i-1,k+1/2}}{\Delta
x}$
$\displaystyle-\frac{(\lambda_{i+1,n-1/2}-\lambda_{i-1,n-1/2})w_{i,n-1/2}}{4\Delta
x}-\frac{(\lambda_{i+1,n+1/2}-\lambda_{i+1,n-1/2})w_{i,n+1/2}}{4\Delta x}$
$\displaystyle+\frac{(w_{i,n-1/2})^{2}}{4}+\frac{(w_{i,n+1/2})^{2}}{4}-\frac{\nu^{4}}{8}\left(-\frac{\rho_{i,n-1/2}-\rho_{i-2,k-1/2}}{8(\rho_{i-1,n-1/2})(\Delta
x)^{2}}-\frac{\rho_{i,n+1/2}-\rho_{i-2,k+1/2}}{8(\rho_{i-1,n+1/2})(\Delta
x)^{2}}\right)$
$\displaystyle-\frac{\nu^{4}}{8}\left(-\frac{\rho_{i,n-1/2}-\rho_{i+2,k-1/2}}{8(\rho_{i+1,n-1/2})(\Delta
x)^{2}}-\frac{\rho_{i,n+1/2}-\rho_{i+2,k+1/2}}{8(\rho_{i+1,n+1/2})(\Delta
x)^{2}}\right)$
$\displaystyle-\frac{\nu^{4}}{16}\left(-\frac{1}{8}\frac{\left(\rho_{i+1,n-1/2}-\rho_{i-1,n-1/2}\right)^{2}}{(\rho_{i,n-1/2})^{2}(\Delta
x)^{2}}-\frac{\rho_{i+1,n-1/2}-\rho_{i-1,n-1/2}}{\rho_{i,k-1/2}(\Delta
x)^{2}}\right)$
$\displaystyle-\frac{\nu^{4}}{8}\left(\frac{1}{2}\frac{\rho_{i+1,n-1/2}-\rho_{i-1,n-1/2}}{(\rho_{i,n-1/2})(\Delta
x)^{2}}+\frac{1}{8}\frac{(\rho_{i+1,n-1/2}-\rho_{i-1,n-1/2})^{2}}{(\rho_{i,n-1/2})^{2}(\Delta
x)^{2}}\right)$
$\displaystyle-\frac{\nu^{4}}{16}\left(-\frac{1}{8}\frac{\left(\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2}\right)^{2}}{(\rho_{i,n+1/2})^{2}(\Delta
x)^{2}}-\frac{\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2}}{\rho_{i,k+1/2}(\Delta
x)^{2}}\right)$
$\displaystyle-\frac{\nu^{4}}{8}\left(\frac{1}{2}\frac{\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2}}{(\rho_{i,n+1/2})(\Delta
x)^{2}}+\frac{1}{8}\frac{(\rho_{i+1,n+1/2}-\rho_{i-1,n+1/2})^{2}}{(\rho_{i,n+1/2})^{2}(\Delta
x)^{2}}\right),$

for $i=1..N_{x}$, $j=1..N_{y}$ and $n=1...N_{t}$. The discretisation above is
multisymplectic, in the sense that it preserves the discrete counterpart of
the multisymplectic form formula (2.12). To do so, first, we need to compute
_the discrete Cartan one-forms_ which are the result of the contributions of
the exterior derivative $d\mathcal{S}_{d}$ from the boundaries. The discrete
Cartan one-forms are computed according to the formula

$\displaystyle\Theta_{d}^{w,\alpha}$
$\displaystyle=\frac{\partial\widehat{L}}{\partial(w)_{i+\alpha_{1},n+\alpha_{2}}}dw_{i+\alpha_{1},n+\alpha_{2}},$
(4.4) $\displaystyle\Theta_{d}^{\rho,\alpha}$
$\displaystyle=\frac{\partial\widehat{L}}{\partial\rho_{i+\alpha_{1},n+\alpha_{2}}}d\rho_{i+\alpha_{1},n+\alpha_{2}},$
$\displaystyle\text{and}\quad\Theta_{d}^{\lambda,\alpha}$
$\displaystyle=\frac{\partial\widehat{L}}{\partial\lambda_{i+\alpha_{1},n+\alpha_{2}}}d\lambda_{i+\alpha_{1},n+\alpha_{2}},$

where $\alpha$ is a multi-index $\alpha=(\alpha_{1},\dots,\alpha_{N+1})$. The
one-forms in (4.4) are for a single $\Delta_{i,n}$ and for the contribution of
the boundaries of the whole action $d\mathcal{S}_{d}$, we need to sum all
these one-forms for all $\Delta_{i,n}$, $i=1,\dots,N_{x}$ and
$n=1,\dots,N_{t}$. As with taking variations of an action functional, as seen
in 2.1, the boundary terms contribute to $d\mathcal{S}_{d}$ and when
restricted to the vector field $V$ are

$\displaystyle\frac{\partial\widehat{L}}{\partial\phi_{i,n}}\cdot(V_{i}^{n})=(j^{1}(\phi_{d}))^{\ast}\left(j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Theta_{d}^{z,\alpha}\right),$

and with that, the total contribution of the one-forms to $d\mathcal{S}_{d}$
is

$\displaystyle\mathrm{d}\mathcal{S}_{d}\cdot
V=\sum\limits_{\mathbb{B}}\sum\limits_{\alpha}(j^{1}(\phi_{d}))^{\ast}\left(j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Theta_{d}^{z,\alpha}\right).$

Taking the exterior derivative of the of $d\mathcal{S}_{d}\cdot V$ and
restricting it to vector field $W$ yields $d^{2}\mathcal{S}_{d}=0$ and that
tells us the sum is

$\displaystyle\sum\limits_{\mathbb{B}}\sum\limits_{\alpha}(j^{1}(\phi_{d}))^{\ast}\left(j^{1}(V)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}j^{1}(W)\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\Omega_{d}^{z,\alpha}\right)=0,$
(4.5)

where $\Omega_{d}^{z,\alpha}=-d\Theta_{d}^{z,\alpha}$. In fact, (4.5) is
analogous to (2.12) and it is known as _the discrete multisymplectic form
formula_ [86]. The importance of this is that when the discrete
multisymplectic form formula is satisfied it means the discretisation scheme
is multisymplectic. With the multisymplectic integrator at hand, the training
problem (3.7) can now be stated as:

###### Problem 4.1.

Let $(Q,\mathcal{F},\mathbb{P})$ be a probability space and the stochastic
control system $\Sigma_{\rho}=(\mathrm{Dens}(Q),\nabla,\mathfrak{X}(Q))$. Let
$\\{(\rho^{(k)}_{i,j,0},c^{(k)})\\}$ be a training set with
$\rho^{(k)}_{i,j,0}\in\mathrm{Dens}(Q)$ and $c^{(k)}\in\mathcal{C}$ for
$k=0,\dots,N$ be . The training of the stochastic ResNet (3.1) is done by
finding a set $(w^{1}_{i,j,n},w^{2}_{i,j,n},\lambda_{i,j,n},\rho_{i,j,n})\in
J^{1}(E)$, such that

$\mathcal{S}_{d,\Sigma_{\rho}}((w^{1})^{\ast}_{i,j,n},(w^{2})^{\ast}_{i,j,n},\rho^{\ast}_{i,j,n})<\mathcal{S}_{d,\Sigma_{\rho}}(w^{1}_{i,j,n},w^{2}_{i,j,n},\rho_{i,j,n}),\quad\forall(w^{1}_{i,j,n},w^{2}_{i,j,n},\rho_{i,j,n})\in
J^{1}(E),$

where the cost functional is defined by

$\displaystyle\mathcal{S}_{d,\Sigma_{\rho}}$
$\displaystyle=\sum\limits_{(i,j,n)\in\mathbb{B}}\widehat{L}\left(\Delta_{i,j,n},w^{1}_{i,j,n},w^{2}_{i,j,n},\lambda_{i,j,n},\rho_{i,j,n}\right)\Delta
x^{1}\Delta x^{2}\Delta t$ (4.6)
$\displaystyle+\sum\limits_{k=0}^{N}\Phi(\pi(\rho_{i,j,N_{t}}),c^{(k)}),$

where $\widehat{L}$ is the cost Lagrangian defined in (4.1) and also satisfies
the discrete Euler–Lagrange equation (4.2).

###  4.1 Applications of multisymplectic integrators

To demonstrate the application of the fluid model for machine learning, here
we apply our hydrodynamic formulation and multisymplectic variational
integrator to solve machine learning problems: density matching and generating
new samples using autoencoder. The direct involvement of probability density
function in the hydrodynamic framework for deep learning naturally lends
itself to a branch of machine learning known as generative models. The main
aim of such models is to learn the joint probability distribution from data
and generate new data points, e.g., generate new samples of a dataset. Study
of such models has gained a considerable interest in recent years with the
introduction of generative adversarial networks (GAN) [45], normalizing flow
based generative model [110, 49, 124, 104], variational autoencoders [5], and
restricted Boltzmann machines [96] to name a few. Each method implements
generative models in a different way. The closest approach to the one
presented here is the normalizing flow approach, especially its variant, deep
diffeomorphic normalizing flows [112].

####  4.1.1 Density matching

The first problem we consider is approximating posterior distributions, which
need not to be Gaussian. One solution approach that gained traction was
initiated in [110], where a discrete dynamical system was used to push-forward
data sampled from a Gaussian distribution to a target distribution. The push-
forward, implemented by a neural network, is applied successively $K$ times.
Here, we begin with the continuous dynamical system

$\displaystyle(\dot{x}^{1},\dot{x}^{2})=\left(w^{1}(x^{1},x^{2},t),w^{2}(x^{1},x^{2},t)\right),$
(4.7)

where $(x^{1},x^{2})$ are data samples, $(w^{1},w^{2}):Q\times\mathbb{R}\to
TQ$ is the vector field whose flow, $F_{(w^{1},w^{2})}$, is the transformation
we are aiming to find. In other words, the flow of $(w^{1},w^{2})$ maps the
data $(x^{1},x^{2})\in Q$ to $(z^{1},z^{2})\in\widetilde{Q}$

$\displaystyle(z^{1},z^{2})=\left(F_{(w^{1},w^{2})}\right)_{\ast}((x^{1},x^{2})),$
(4.8)

where both $Q,\widetilde{Q}\subset\mathbb{R}^{2}$ are smooth manifolds. There
are many variations on the theme of normalizing flows, where $(w^{1},w^{2})$
can be fixed to be a specific type, for example, it can be a constant
function, whose flow is a linear function and this type is known as the planar
flow, as in the case of the original paper [110]. Another is to impose
$(w^{1},w^{2})$ a diffeomorphism, namely a deep diffeomorphic normalizing
flows in [112]. Another interesting choice of $(w^{1},w^{2})$ is choosing it
as a potential function, which was introduced by [126], thus the flow also
solves of the Monge–Ampère equation. The latter is connected to the approach
of [10], where they showed that Monge–Ampère is also the solution of optimal
transport problem [10]. The main difference between the approach here and the
aforementioned ones, is that here the flow $F_{(w^{1},w^{2})}$ is computed
directly on a fixed Eulerian grid, whereas in [110, 126] deep neural networks
have been employed to approximate flow using Lagrangian coordinates.

The idea behind finding $F_{(w^{1},w^{2})}$ is to enable us to learn the joint
probability density function from sampled data and once the transformation is
obtained, new data samples can be generated by, first, generating samples from
distributions on $\mathrm{Dens}(\widetilde{Q})$ and then transforming them to
new samples from $\mathrm{Dens}({Q})$. This is the idea of normalizing flows
and it helps in implementing generative models. In the probability density
formulation, the problem becomes what is known as the _Monge–Kantorovich
optimal transport problem_. Given two density functions $\rho_{0}\geq 0$ and
$\rho_{T}\geq 0$, and optimal transport concerns finding a smooth map, $\eta$,
such that $(\eta)_{\ast}\rho_{t}=\rho_{0}$, and at the same time the map $g$
minimises the distance between the two prescribed densities. The distance in
this case is the $p$-Wasserstein distance and it is defined by

$\displaystyle
W_{p}(\rho_{0},\rho_{T})=\inf\limits_{g}\int|g(x^{1},x^{2})-(x^{1},x^{2})|^{p}\rho_{0}(x^{1},x^{2})dx^{1}dx^{2},$
(4.9)

where $|\cdot|^{p}$ is the $L^{p}$ norm on $\mathbb{R}^{d}$. In [10], it was
shown that this problem is equivalent to the following fluid dynamics problem

$\displaystyle\mathcal{S}$
$\displaystyle=\int_{0}^{T}\int_{\mathbb{R}^{d}}\rho(t,x^{1},x^{2})|(w^{1},w^{2})|^{2}+\lambda\left(\partial_{t}\rho+\nabla\cdot(\rho(w^{1},w^{2}))\right)dx^{1}dx^{2}\,dt$
(4.10)
$\displaystyle+\int_{0}^{T}\Pi\left(\int\rho(t,x^{1},x^{2})dx^{1}dx^{2}-1\right),$

such that $\rho(0,x^{1},x^{2})=\rho_{0}(x^{1},x^{2})$ and
$\rho(T,x^{1},x^{2})=\rho_{T}(x^{1},x^{2})$. The term $\Pi$ is the Lagrange
multiplier that enforces the incompressibility condition, and in hydrodynamics
terms it acts as the pressure force. The optimality condition gives
$(w^{1},w^{2})=\nabla\lambda$ and the Lagrange multiplier satisfies the
Hamilton–Jacobi equation. Taking gradient of the Hamilton–Jacobi equation
gives an equation for $u$ and it is _incompressible Euler’s fluid equation_

$\displaystyle\partial_{t}(w^{1},w^{2})+\left((w^{1},w^{2})\cdot\nabla\right)(w^{1},w^{2})=-\nabla\Pi,\quad\nabla\cdot(w^{1},w^{2})=0.$
(4.11)

As pointed out in remark 3.5, the system 3.7, when $\sigma\to 0$ we recover
the optimal transport of Benamou and Brenier [10] and it can be viewed as its
regularisation. Given the initial probability density
$\rho_{0}(x^{1},x^{2})=\mathcal{N}(0,\mathbf{I})\in\mathrm{Dens}\left(\mathbb{R}^{2}\right)$
and the target density $\rho_{T}=e^{-U(x^{1},x^{2})}$. To test our method,
$U(x^{1},x^{2})$ is chosen as

$U(x^{1},x^{2})=\begin{cases}\frac{1}{2}\left(\frac{\|x^{1}\|-2}{0.4}\right)^{2}-\mathrm{ln}\left(e^{-\frac{1}{2}\left(\frac{x^{1}-2}{0.4}\right)^{2}}+e^{-\frac{1}{2}\left(\frac{x^{1}+2}{0.6}\right)^{2}}\right),\\\
\frac{1}{2}\left(\frac{x^{2}-\mathrm{sin}\left(\frac{2\pi
x^{1}}{4}\right)}{0.4}\right)^{2},\\\
-\mathrm{ln}\left(e^{-\frac{1}{2}\left(\frac{x^{2}-\mathrm{sin}\left(\frac{2\pi
x^{1}}{4}\right)}{0.35}\right)^{2}}+e^{-\frac{1}{2}\left(\frac{x^{2}-\mathrm{sin}\left(\frac{2\pi
x^{1}}{4}\right)+3e^{-\frac{1}{2}\left(\frac{(x^{1}-1)}{0.6}\right)^{2}}}{0.35}\right)^{2}}\right),\\\
-\mathrm{ln}\left(e^{-\frac{1}{2}\left(\frac{x^{2}-\mathrm{sin}\left(\frac{2\pi
x^{1}}{4}\right)}{0.4}\right)^{2}}+e^{-\frac{1}{2}\left(\frac{x^{2}-\mathrm{sin}\left(\frac{2\pi
x^{1}}{4}\right)+3\left(1+e^{-\frac{x^{1}-1}{0.3}}\right)^{-1}}{0.35}\right)^{2}}\right),\end{cases}$
(4.12)

The aim is to find a vector field $\theta$ that would steer $\rho$ form
$\rho_{0}$ to $\rho_{T}$ while minimising the cost functional

$\displaystyle\mathcal{S}=\Phi(\rho(x^{1},x^{2},T),\rho_{T}(x^{1},x^{2}))+\frac{1}{2}\int_{0}^{T}\int_{\mathbb{R}^{2}}\rho(x^{1},x^{2},t)\|\theta(x^{1},x^{2},t)\|^{2}\,dx^{1}dx^{2}dt.$
(4.13)

This test was introduced in [110] and since then it has become the de facto
benchmark. Here, density matching is done using $\Sigma_{\rho}$ control system
and the training optimal control problem is solved using a variational
integrator described in section 4. For this particular problem, centred
difference discretisation, whose Lagrangian is given by (4.1), is used. The
end time $T$ was chosen to be $1$ and $16$ discretisation points in space are
used. For $x^{1}$ and $x^{2}$-spaces, have $32$ grid points for each space. As
the optimality condition is a system of algebraic equations (4.2),it is solved
using a newton-type solver at each time step. The result is shown in figure 2.

Figure 2:  The problem 4.1 is solved numerically, with potentials in (4.12)
are chosen for target densities $\rho_{T}$ and in each case, the initial
density is specified as a Gaussian normal distribution whose mean is zero and
its covariance matrix is the identity matrix. Each row corresponds to a
density in (4.12). The left most column is a contour plot of the target
density. The second column from the left represents the numerical solution of
problem 4.1 after one time step. The third column from the left is the
solution after seven time iterations. Finally, the rightmost column is the
solution at the final iteration.

####  4.1.2 Generative model based on multisymplectic integrators

The concept of autoencoders is an interesting one, which are artificial neural
networks [71, 44]. The training for such networks is done in an unsupervised
manner by performing input data mimicry, that is learn coded representation,
in terms of latent state/features, of the data and from it reconstruct the
inputs. Let $\pi_{Q\mathrm{Z}}:Q\to\mathrm{Z}$ be a fibre bundle and we also
define the map $\Gamma_{\mathrm{Z}Q}:\mathrm{Z}\to Q$. Here, an autoencoder
consists of an encoder $\pi_{Q\mathrm{Z}}$ and a decoder
$\Gamma_{\mathrm{Z}Q}$, where the base manifold $\mathrm{Z}$ is the space of
latent state/features. This renders autoencoders suitable for data compression
and dimensionality reduction [51] applications.

Looking at autoencoders from a probabilistic view point, where the latent
states are sampled from a probability distribution, it allows us to use them
as a generative model. Such models learn the representation of the dataset and
allow the generation of new data. Such a class of autoencoders is known as
_variational autoencoders_ [68]. To generate new data, a random variable is
sampled from the distribution $\rho_{0}\in\mathrm{Dens}(\mathsf{Z})$ and then
feeding into the decoder, $\Gamma_{\mathrm{Z}Q}(z)$, to obtain the
reconstructed data. Variational autoencoders are often trained such that
latent variables $z$ are sampled from $\mathrm{Dens}(\mathsf{Z})$ would have
either a normal or uniform probability distribution to simplify sampling from
a distribution. This is done by introducing Kullback-–Leibler regularising
term to the cost function. The downside of forcing the distribution of the
parameter space to be an elementary one renders reconstruction less accurate.
In recent years, there has been an interest in the using optimal transport in
conjunction with autoencoders to develop an architecture as a substitute or a
complement to generative adversarial networks [117, 111, 114, 83]. For
example, in [18] the problem of generative modelling was cast as an optimal
transport problem and in [42] the measure and cost function of optimal
transport, solved using Sinkhorn algorithm, are approximated by neural
networks. Whereas in [125], single-cell images are mapped to a latent space
and then the cell trajectories are predicted using optimal transport. The use
of normalizing flows as a generative model is has been studied in works such
as [105, 67] in what is known as _flow-based generative models_ , however such
models do not incorporate an autoencoder structure. The use of autoencoder
with normalizing flows, replacing optimal transport, is used for the sampling
step generative models and this have been studied in [63, 95].

Inspired by [4] and flow-based generative models, the idea is to not to force
the variational autoencoder to have the distribution of the latent states as
an elementary distribution, instead train the autoencoder to solely perform
dimensionality reduction. Here, a control system $\Sigma_{\rho}$ is used to
map $\rho_{0}\in\mathrm{Dens}(\mathsf{Z})$, where it is the probability
density function of the latent states, to a normal distribution,
$\rho_{e}\simeq\mathcal{N}(0,\mathbf{I})$. In [4] that was done by solving a
Monge–Kantorovich problem. Random variables are sampled from a normal
distribution and then pushed-forward by the flow of solution of the
Monge–Kantorovich and finally they are fed into the decoder. However, solving
optimal transport problems is challenging and instead we use the
multisymplectic variational integrator for that. The density of the space of
latent states $\rho_{0}\in\mathrm{Dens}(\mathsf{Z})$ is estimated using the
presented method here. In the second step of the reconstruction, a latent
variable $z$ is sampled from a normal distribution and then transported using
diffeomorphisms to a random variable $\widetilde{z}$ whose probability
distribution has the density $\rho_{0}$. The diffeomorphism here is an
integral curve of the vector field $u\in\mathfrak{X}(Q)$ and it is learned by
solving problem 3.7 using a multisymplectic integrator. Once training is over,
new data can be generated by sampling from $\mathcal{N}(0,\mathbb{I})$ and
then mapping it to $\mathsf{Z}$. The transformed data is then fed into the
decoder, and the result of one experiment can be seen in figure 3. The
construction presented here can be viewed as an instance of pullback bundle.

Figure 3:  New handwritten digits are generated using the autoencoder with a
continuous ResNet trained using multisymplectic optimal control. Here, an
autoencoder is trained to reconstruct images from MNIST dataset [77], where
$Q$ is $\mathbb{R}^{28\times 28}$ and each element is an image of a
handwritten digit. The space of latent variable is chosen as
$\mathsf{Z}\in\mathbb{R}^{2}$. A set of points that form a square lattice
defined on $\mathbb{R}^{2}$ are mapped to $\mathsf{Z}$ using the flow of the
deep neural network (3.1), whose parameters, $\theta=u$, are obtained by
solving problem 4.1. The points in the transformed lattice are then mapped to
$Q$ using the decoder.

##  5 Backward error analysis of multisymplectic integrator

Different numerical schemes can be applied to solve the optimal control
problem for training, but as mentioned, schemes based on multisymplectic
geometry are chosen in order to obtain an approximation that is closely
related to the exact system. In numerical analysis, there are mainly two ways
of analysing the qualitative performance of the scheme; forward analysis, and
backward error analysis. In forward error analysis, the solution being a curve
is compared to a known integral curve of the system. Many error measurements
in machine learning are based on this and it requires knowing the solution or
availability of the dataset. Backward error analysis is concerned computing a
differential equation whose exact solution is the numerical solution and this
equation is known as _modified equation_ and it offers a paradigm to conduct a
study that enables direct comparison between a continuous multisymplectic
system and its discrete counterpart [107].

One of the main favourable properties of multisymplectic integrators is that
their modified counterparts are also multisymplectic, as demonstrated in [92,
61]. Once the modified equation is known, we obtain time evolution of the
network that explicitly includes parameters such as the number of hidden
layers and nodes per layer. This is valuable, as it provides a way to study
the qualitative behaviour of the solution in response to the change in the
network’s parameters.

In order to compute the modified equations, the solution
$\phi=(w^{1},w^{2},\lambda,\rho)\in J^{1}(E)$ is assumed to be a $4$-tuple of
smooth functions. As the multisymplectic integrator uses points in
$J^{1}(\mathbb{E})$ to approximate the solution, their Taylor expansion around
the centre of $\Delta_{i,j}^{n}$, $\bar{x}$, on the base manifold $Q$ allows
to move from discrete points on $\mathbb{E}$ to a section on $\mathbb{E}$. To
illustrate that, the Taylor expansion of $\phi_{i+1/2,j+1/2,n+1}$ and
$\phi_{i+1/2,j+1/2,n}$ yields:

$\displaystyle\phi_{i+1/2,j+1/2,n+1}$
$\displaystyle=\phi(\bar{x})+\frac{(\Delta
t)}{2}\partial_{t}\phi(\bar{x})+\frac{(\Delta
t)^{2}}{8}\partial^{2}_{t}\phi(\bar{x})+\frac{(\Delta
t)^{3}}{48}\partial_{ttt}\phi(\bar{x})+\mathcal{O}((\Delta t)^{4}),$
$\displaystyle\phi_{i+1/2,j+1/2,n}$ $\displaystyle=\phi(\bar{x})-\frac{(\Delta
t)}{2}\partial_{t}\phi(\bar{x})+\frac{(\Delta
t)^{2}}{8}\partial^{2}_{t}\phi(\bar{x})-\frac{(\Delta
t)^{3}}{48}\partial_{ttt}\phi(\bar{x})+\mathcal{O}((\Delta t)^{4}).$

In addition, the approximation of the section over the $\bar{x}$, which is

$\displaystyle\bar{\phi}$
$\displaystyle=\frac{1}{8}\left(\phi_{i+1,j+1,n+1}+\phi_{i,j+1,n+1}+\phi_{i+1,j,n+1}+\phi_{i,j,n+1}\right.$
(5.1)
$\displaystyle\left.+\phi_{i+1,j+1,n}+\phi_{i,j+1,n}+\phi_{i+1,j,n}+\phi_{i,j,n}\right)$
$\displaystyle=\phi(\bar{x})+\frac{1}{8}\left((\Delta
t)^{2}\partial^{2}_{t}\phi+(\Delta x^{1})^{2}\partial^{2}_{x^{1}}\phi+(\Delta
x^{2})^{2}\partial^{2}_{x^{2}}\phi\right)$
$\displaystyle+\mathcal{O}\left(|(\Delta t)^{3}|+|(\Delta x^{1})^{3}|+|(\Delta
x^{2})^{3}|\right).$

The time discretisation using scheme in (4.1) is

$\displaystyle\frac{\phi_{i+1/2,j+1/2,n+1}-\phi_{i+1/2,j+1/2,n}}{(\Delta t)}$
$\displaystyle=\partial_{t}\phi+\frac{(\Delta
t)^{2}}{24}\partial_{ttt}\phi+\mathcal{O}((\Delta t)^{4}).$ (5.2)

Likewise, the Taylor expansion of spatial discretisations are

$\displaystyle\frac{\phi_{i+1,j+1/2,n+1/2}-\phi_{i,j+1/2,n+1/2}}{(\Delta
x^{1})}$ $\displaystyle=\partial_{x^{1}}\phi+\frac{(\Delta
x^{1})^{2}}{24}\partial_{x^{1}x^{1}x^{1}}\phi+\mathcal{O}((\Delta
x^{1})^{4}),$ (5.3)
$\displaystyle\frac{\phi_{i+1/2,j+1,n+1/2}-\phi_{i+1/2,j,n+1/2}}{(\Delta
x^{2})}$ $\displaystyle=\partial_{x^{2}}\phi+\frac{(\Delta
x^{2})^{2}}{24}\partial_{x^{2}x^{2}x^{2}}\phi+\mathcal{O}((\Delta
x^{2})^{4}).$

Substituting the Taylor expansion of the finite difference relations into the
stationary variations of the action whose Lagrangian is give by (4.1) and
rearranging the terms of the discretised Fokker–Planck equation, the system of
algebraic equations are then become the following system of partial
differential equations:

$\displaystyle\delta\lambda$ $\displaystyle:\quad
0=\partial_{t}\rho+\partial_{x^{1}}(\rho\phi(w^{1}))+\partial_{x^{2}}(\rho\phi(w^{2}))+\frac{(\Delta
x^{1})^{2}}{24}\partial_{x^{1}x^{1}x^{1}}(\rho\phi(w^{1}))$ (5.4)
$\displaystyle+\frac{(\Delta
x^{2})^{2}}{24}\partial_{x^{2}x^{2}x^{2}}(\rho\phi(w^{2}))-\frac{(\Delta
t)^{2}}{24}\partial_{ttx^{1}}(\rho\phi(w^{1}))-\frac{(\Delta
t)^{2}}{24}\partial_{ttx^{2}}(\rho\phi(w^{2})),$ $\displaystyle\delta\rho$
$\displaystyle:\quad
0=\partial_{t}\lambda+\phi(w^{1})\partial_{x^{1}}\lambda+\phi(w^{2})\partial_{x^{2}}\lambda+\frac{(\Delta
x^{1})^{2}}{24}\phi(w^{1})\partial_{x^{1}x^{1}x^{1}}\lambda$
$\displaystyle+\frac{(\Delta
x^{2})^{2}}{24}\phi(w^{2})\partial_{x^{2}x^{2}x^{2}}\lambda-\frac{(\Delta
x^{1})^{2}}{24}\phi(w^{1})\partial_{x^{1}tt}\lambda-\frac{(\Delta
x^{1})^{2}}{24}\phi(w^{2})\partial_{x^{2}tt}\lambda-\frac{\delta\ell_{m}}{\delta\rho},$
$\displaystyle\delta w^{1}$ $\displaystyle:\quad
0=\frac{\delta\ell_{m}}{\delta
w^{1}}-\rho\partial_{x^{1}}\lambda-\frac{(\Delta
x^{1})^{2}}{24}\rho\partial_{x^{1}x^{1}x^{1}}\lambda+\frac{(\Delta
t)^{2}}{24}\rho\partial_{x^{1}tt}\lambda,$ $\displaystyle\delta w^{2}$
$\displaystyle:\quad 0=\frac{\delta\ell_{m}}{\delta
w^{2}}-\rho\partial_{x^{2}}\lambda-\frac{(\Delta
x^{2})^{2}}{24}\rho\partial_{x^{2}x^{2}x^{2}}\lambda+\frac{(\Delta
t)^{2}}{24}\rho\partial_{x^{2}tt}\lambda.$

The above system of partial differential equations satisfies the stationary
variations of the the action functional

$\displaystyle\mathcal{S}=\int_{0}^{t}\int_{Q}L_{m}\,d^{n}x\,dt,$ (5.5)

i.e. $0=\delta\mathcal{S}$, and with a modified Lagrangian given by

$\displaystyle L_{m}$ $\displaystyle=L+\frac{(\Delta
x^{1})^{2}}{24}L_{1}+\frac{(\Delta x^{2})^{2}}{24}L_{2}+\frac{(\Delta
t)^{2}}{24}L_{3},$ (5.6)

where the terms of the Lagrangian are given by

$\displaystyle L$
$\displaystyle=+\frac{1}{2}\rho(w^{1})^{2}+\frac{1}{2}\rho(w^{2})^{2}+\frac{\nu^{2}}{2}w^{1}\partial_{x^{1}}\rho+\frac{\nu^{2}}{2}w^{2}\partial_{x^{2}}\rho+\frac{\nu^{4}\left(\partial_{x^{1}}\rho\right)^{2}}{8\rho+(\Delta
x^{1})^{2}\partial_{x^{1}x^{1}}\rho+(\Delta
x^{2})^{2}\partial_{x^{2}x^{2}}\rho+(\Delta t)^{2}\partial_{tt}\rho}$
$\displaystyle+\frac{\nu^{4}\left(\partial_{x^{2}}\rho\right)^{2}}{8\rho+(\Delta
x^{1})^{2}\partial_{x^{1}x^{1}}\rho+(\Delta
x^{2})^{2}\partial_{x^{2}x^{2}}\rho+(\Delta
t)^{2}\partial_{tt}\rho}+\lambda\left(\partial_{t}\rho+\partial_{x^{1}}(\rho(w^{1}))+\partial_{x^{2}}(\rho(w^{2}))\right),$
$\displaystyle L_{1}$ $\displaystyle=\frac{(\Delta x^{1})^{2}}{4}\rho
w^{1}\partial_{x^{1}x^{1}}w^{1}+\frac{(\Delta x^{1})^{2}}{4}\rho
w^{2}\partial_{x^{1}x^{1}}w^{2}+\frac{(\Delta
x^{1})^{2}}{8}\rho\left(\partial_{x^{1}}w^{1}\right)^{2}+\frac{(\Delta
x^{1})^{2}}{8}\rho\left(\partial_{x^{1}}w^{2}\right)^{2}$
$\displaystyle+\frac{\nu^{2}(\Delta
x^{1})^{2}}{48}w^{1}\partial_{x^{1}x^{1}x^{1}}\rho+\frac{\nu^{2}(\Delta
x^{1})^{2}}{16}(\partial_{x^{1}}\rho)\partial_{x^{1}x^{1}}w^{1}+\frac{\nu^{2}(\Delta
x^{1})^{2}}{16}(\partial_{x^{2}}\rho)\partial_{x^{1}x^{1}}w^{2}+\lambda\left(\partial_{x^{1}x^{1}x^{1}}(w^{1}\rho)\right),$
$\displaystyle L_{2}$ $\displaystyle=\frac{(\Delta x^{2})^{2}}{4}\rho
w^{1}\partial_{x^{2}x^{2}}w^{1}+\frac{(\Delta x^{2})^{2}}{4}\rho
w^{2}\partial_{x^{2}x^{2}}w^{2}+\frac{(\Delta
x^{2})^{2}}{8}\rho\left(\partial_{x^{2}}w^{1}\right)^{2}+\frac{(\Delta
x^{2})^{2}}{8}\rho\left(\partial_{x^{2}}w^{2}\right)^{2}$
$\displaystyle+\frac{\nu^{2}(\Delta
x^{2})^{2}}{48}w^{2}\partial_{x^{2}x^{2}x^{2}}\rho+\frac{\nu^{2}(\Delta
x^{2})^{2}}{16}(\partial_{x^{1}}\rho)\partial_{x^{2}x^{2}}w^{1}+\frac{\nu^{2}(\Delta
x^{2})^{2}}{16}(\partial_{x^{2}}\rho)\partial_{x^{2}x^{2}}w^{2}+\lambda\left(\partial_{x^{2}x^{2}x^{2}}(\rho(w^{2}))\right),$
$\displaystyle L_{3}$ $\displaystyle=\frac{(\Delta t)^{2}}{4}\rho
w^{1}\partial_{tt}w^{1}+\frac{(\Delta t)^{2}}{4}\rho
w^{2}\partial_{tt}w^{2}+\frac{(\Delta
t)^{2}}{8}\rho\left(\partial_{t}w^{1}\right)^{2}+\frac{(\Delta
t)^{2}}{8}\rho\left(\partial_{t}w^{2}\right)^{2}-\frac{\nu^{2}(\Delta
t)^{2}}{48}w^{1}\partial_{x^{1}tt}\rho$ $\displaystyle-\frac{\nu^{2}(\Delta
t)^{2}}{48}w^{2}\partial_{x^{2}tt}\rho+\frac{\nu^{2}(\Delta
t)^{2}}{16}(\partial_{x^{1}}\rho)\partial_{tt}w^{1}+\frac{\nu^{2}(\Delta
t)^{2}}{16}(\partial_{x^{2}}\rho)\partial_{tt}w^{2}-\lambda\left(\partial_{x^{1}tt}(\rho(w^{1}))+\partial_{x^{2}tt}(\rho(w^{2}))\right).$

For the modified equation, the cost Lagrangian $\ell_{m}$, that appears in
(5.4), is

$\displaystyle\ell_{m}$
$\displaystyle=+\frac{1}{2}\rho(w^{1})^{2}+\frac{1}{2}\rho(w^{2})^{2}+\frac{\nu^{2}w^{1}\partial_{x^{1}}\rho}{2}+\frac{\nu^{2}w^{2}\partial_{x^{2}}\rho}{2}$
(5.7)
$\displaystyle+\frac{\nu^{4}\left(\partial_{x^{1}}\rho\right)^{2}}{8\rho+(\Delta
x^{1})^{2}\partial_{x^{1}x^{1}}\rho+(\Delta
x^{2})^{2}\partial_{x^{2}x^{2}}\rho+(\Delta
t)^{2}\partial_{tt}\rho}+\frac{(\Delta x^{1})^{2}}{4}\rho
w^{1}\partial_{x^{1}x^{1}}w^{1}$
$\displaystyle+\frac{\nu^{4}\left(\partial_{x^{2}}\rho\right)^{2}}{8\rho+(\Delta
x^{1})^{2}\partial_{x^{1}x^{1}}\rho+(\Delta
x^{2})^{2}\partial_{x^{2}x^{2}}\rho+(\Delta
t)^{2}\partial_{tt}\rho}+\frac{(\Delta x^{1})^{2}}{4}\rho
w^{2}\partial_{x^{1}x^{1}}w^{2}$ $\displaystyle+\frac{(\Delta
x^{1})^{2}}{8}\rho\left(\partial_{x^{1}}w^{1}\right)^{2}+\frac{(\Delta
x^{1})^{2}}{8}\rho\left(\partial_{x^{1}}w^{2}\right)^{2}+\frac{\nu^{2}(\Delta
x^{1})^{2}}{48}w^{1}\partial_{x^{1}x^{1}x^{1}}\rho$
$\displaystyle+\frac{\nu^{2}(\Delta
x^{1})^{2}}{16}(\partial_{x^{1}}\rho)\partial_{x^{1}x^{1}}w^{1}+\frac{\nu^{2}(\Delta
x^{1})^{2}}{16}(\partial_{x^{2}}\rho)\partial_{x^{1}x^{1}}w^{2}+\frac{(\Delta
x^{2})^{2}}{4}\rho w^{1}\partial_{x^{2}x^{2}}w^{1}$
$\displaystyle+\frac{(\Delta x^{2})^{2}}{4}\rho
w^{2}\partial_{x^{2}x^{2}}w^{2}+\frac{(\Delta
x^{2})^{2}}{8}\rho\left(\partial_{x^{2}}w^{1}\right)^{2}+\frac{(\Delta
x^{2})^{2}}{8}\rho\left(\partial_{x^{2}}w^{2}\right)^{2}$
$\displaystyle+\frac{\nu^{2}(\Delta
x^{2})^{2}}{48}w^{2}\partial_{x^{2}x^{2}x^{2}}\rho+\frac{\nu^{2}(\Delta
x^{2})^{2}}{16}(\partial_{x^{1}}\rho)\partial_{x^{2}x^{2}}w^{1}+\frac{\nu^{2}(\Delta
x^{2})^{2}}{16}(\partial_{x^{2}}\rho)\partial_{x^{2}x^{2}}w^{2}$
$\displaystyle+\frac{(\Delta t)^{2}}{4}\rho
w^{1}\partial_{tt}w^{1}+\frac{(\Delta t)^{2}}{4}\rho
w^{2}\partial_{tt}w^{2}+\frac{(\Delta
t)^{2}}{8}\rho\left(\partial_{t}w^{1}\right)^{2}$ $\displaystyle+\frac{(\Delta
t)^{2}}{8}\rho\left(\partial_{t}w^{2}\right)^{2}-\frac{\nu^{2}(\Delta
t)^{2}}{48}w^{1}\partial_{x^{1}tt}\rho-\frac{\nu^{2}(\Delta
t)^{2}}{48}w^{2}\partial_{x^{2}tt}\rho+\frac{\nu^{2}(\Delta
t)^{2}}{16}(\partial_{x^{1}}\rho)\partial_{tt}w^{1}$
$\displaystyle+\frac{\nu^{2}(\Delta
t)^{2}}{16}(\partial_{x^{2}}\rho)\partial_{tt}w^{2}$

Unlike the problem 3.7, the modified system belongs to the third jet
prolongation of the Pontryagin’s bundle
$E:=T^{\ast}\mathrm{Dens}(Q)\times\mathfrak{X}(Q)$, i.e $j^{3}(E)$ which is a
manifold with local coordinates

$\displaystyle j^{3}(z)$
$\displaystyle=\left(x^{1},x^{2},t,z,z_{x^{1}x^{1}},z_{x^{1}x^{2}},z_{t},z_{x^{2}x^{1}},z_{tt},z_{x^{2}x^{2}},z_{x^{1}t},z_{tx^{1}},z_{x^{2}x^{1}},z_{x^{1}x^{2}},z_{x^{2}t},z_{tx^{2}},\right.$
$\displaystyle\left.z_{x^{1}x^{1}x^{1}},z_{x^{1}x^{1}x^{2}},z_{x^{1}x^{1}t},z_{x^{1}x^{2}x^{1}},z_{x^{1}x^{2}x^{2}},z_{x^{1}x^{2}t},z_{x^{1}tx^{1}},z_{x^{1}tx^{2}},z_{x^{1}tt},z_{x^{2}x^{1}x^{1}},z_{x^{2}x^{1}x^{2}},\right.$
$\displaystyle\left.z_{x^{2}x^{1}t},z_{x^{2}x^{2}x^{1}},z_{x^{2}x^{2}x^{2}},z_{x^{2}x^{2}t},z_{x^{2}tx^{1}},z_{x^{2}tx^{2}},z_{x^{2}tt},z_{tx^{1}x^{1}},z_{tx^{1}x^{2}},z_{tx^{1}t},z_{tx^{2}x^{1}},z_{tx^{2}x^{2}},\right.$
$\displaystyle\left.z_{tx^{2}t},z_{ttx^{1}},z_{ttx^{2}},z_{ttt}\right),$

where the section $z(x^{1},t)$ is a shorthand notation for the tuple

$\displaystyle(w^{1}(x^{1},x^{2},t),w^{2}(x^{1},x^{2},t),\lambda(x^{1},x^{2},t),\rho(x^{1},x^{2},t)).$

With that the modified Lagrangian density
$\widetilde{\mathcal{L}}:j^{3}(z)\to\Lambda^{n+1}(E)$ described by

$\displaystyle\widetilde{\mathcal{L}}\left(j^{3}(z)\right)=L_{m}\,dx^{1}\wedge
dx^{2}\wedge dt,$

where $L_{m}$ is given by (5.6).

The introduction of higher order terms in the modified Lagrangian suggests
that the solution of the modified equations is a section of the third jet
bundle $J^{3}(E)$. The construction of a third jet bundle whose sections are
third jet prolongation of a section $\phi:B\to E$ and they are done by
iterating the first jet prolongation $j^{1}(j^{1}(j^{1}(\phi)))$. In
coordinates, the third jet bundle $j^{3}(\phi)$ is given by

$\displaystyle
x^{\mu}\mapsto\left(x^{\mu},\phi^{A},\frac{\partial\phi^{A}}{\partial
x^{\mu}},\frac{\partial^{2}\phi^{A}}{\partial x^{\mu_{1}}\partial
x^{\mu_{2}}},\frac{\partial^{3}\phi^{A}}{\partial x^{\mu_{1}}\partial
x^{\mu_{2}}\partial x^{\mu_{3}}}\right).$

With that, we consider a modified Lagrangian density
$\mathcal{L}_{m}:J^{3}(E)\to\Lambda^{n+1}(B)$, and in coordinates it is
written as

$\displaystyle\mathcal{L}_{m}=L_{m}\left(x^{\mu},\phi^{A},\frac{\partial\phi^{A}}{\partial
x^{\mu}},\frac{\partial^{2}\phi^{A}}{\partial x^{\mu_{1}}\partial
x^{\mu_{2}}},\frac{\partial^{3}\phi^{A}}{\partial x^{\mu_{1}}\partial
x^{\mu_{2}}\partial x^{\mu_{3}}}\right)dx^{n+1}.$

One quantity that changes with the jet prolongation is the Cartan form on
$\Theta_{\mathcal{L}_{m}}$, which now given by

$\displaystyle\Theta_{\mathcal{L}_{m}}$ $\displaystyle=\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1}}}-D_{\mu_{2}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}}}\right)+D_{\mu_{2}}D_{\mu_{3}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}\mu_{3}}}\right)\right)d\phi^{A}\wedge(\frac{\partial}{\partial\mu_{1}}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}dx^{n+1})$
(5.8) $\displaystyle+\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}}}-D_{\mu_{3}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2},\mu_{3}}}\right)\right)d\phi_{\mu_{1}}^{A}\wedge\left(\frac{\partial}{\partial\mu_{2}}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}dx^{n+1}\right)$
$\displaystyle+\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2},\mu_{3}}}d\phi_{\mu_{1},\mu_{2}}^{A}\wedge\left(\frac{\partial}{\partial\mu_{3}}\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}dx^{n+1}\right)+\left(L-\frac{\partial
L}{\partial\phi^{A}_{\mu_{1}}}\phi^{A}_{\mu_{1}}\right.$
$\displaystyle\left.+D_{\mu_{2}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}}}\right)\phi^{A}_{\mu_{1}}-D_{\mu_{2}}D_{\mu_{3}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}\mu_{3}}}\right)\phi^{A}_{\mu_{1}}-\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2}}}\phi^{A}_{\mu_{1},\mu_{2}}\right.$
$\displaystyle\left.+D_{\mu_{3}}\left(\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2},\mu_{3}}}\right)\phi^{A}_{\mu_{1},\mu_{2}}-\frac{\partial
L}{\partial\phi^{A}_{\mu_{1},\mu_{2},\mu_{3}}}\phi^{A}_{\mu_{1},\mu_{2},\mu_{3}}\right)dx^{n+1}.$

Here, $D_{\mu}$ is the directional derivative of a function $L\in
C^{\infty}(J^{k}(E))$ in the direction $x^{\mu}$ and it is given by

$\displaystyle D_{\mu}\phi=\partial_{\mu}\phi+\frac{\partial
L}{\partial\phi_{\mu_{1}}^{A}}\phi_{\mu}^{A}+\dots+\frac{\partial
L}{\partial\phi_{\mu_{1},\dots,\mu_{k}}^{A}}\phi_{\mu_{1},\dots,\mu_{k}}^{A}.$
(5.9)

Once we have the modified Lagrangian, it is possible to find the conserved
quantities and this is discussed at length in the next section.

##  6 Nonlinear stability of the training algorithm

One of the most attractive features of the hydrodynamics formulation for
training deep neural networks is that allows us to study the nonlinear
stability. In order to do so, we need to find the conserved quantities and the
corresponding conservation laws associated with the hydrodynamic system
(3.21). Applying Noether’s theorem [98], and especially its version for Euler-
Poincaré with advection [32], conserved quantities are identified using vector
fields that generates a symmetry.

First, we define the following manifold $\mathcal{M}$, which is the Hilbert
space closure of the manifold

$\displaystyle\mathcal{M}^{\infty}=\left\\{\phi:U\to
E\quad|\quad\pi_{QE}\circ\phi:U\to Q\right\\},$

where $U$ is a smooth manifold that is a subset of $Q$ and $\pi_{QE}:U\to
U_{Q}$ is a diffeomorphism, where $U_{Q}=\pi_{QE}\circ\phi(U)$. Further, we
define $\mathcal{P}$ to be a manifold defined as

$\displaystyle\mathcal{P}=\left\\{\phi\in\mathcal{M}\quad|\quad
j^{1}(\phi\circ\phi^{-1}_{Q})^{\ast}\left[V\mathbin{\hbox to6.0pt{\vrule
height=0.4pt,width=5.0pt,depth=0.0pt\kern-0.4pt\vrule
height=6.0pt,width=0.4pt,depth=0.0pt\hss}}\Omega_{\mathcal{L}}\right]=0\quad\forall\,V\in
TJ^{1}(E)\right\\}.$

In other words, the elements of $\mathcal{P}$ are the solution of the Euler-
Lagrange equation. Next consider the action functional $\mathcal{S}$ on
$\mathcal{M}$ and consider the infinitesimal transformation given by $\delta
g=\eta\circ g$, where $\eta\in\mathfrak{X}(Q)$ and $g\in\mathrm{Diff}(Q)$ and
from it we can derive the infinitesimal transformation for the vector field
$v$ which is given by

$\displaystyle\delta
w=\dot{\eta}+[w,\eta],\quad\delta\rho=-\pounds_{\eta}\rho.$ (6.1)

Substituting these infinitesimal transformations into $\delta\mathcal{S}=0$,
we obtain

$\displaystyle 0=\int_{0}^{T}\int_{U}\left(\frac{\partial}{\partial
t}\left\langle\frac{\delta\ell}{\delta
w}+\mathrm{ad}^{\ast}_{f(w)}\frac{\delta\ell}{\delta
w}+\rho\diamond\frac{\delta\ell}{\delta\rho}\right),\eta\right\rangle
d^{n}x\,dt+\left.\int_{\partial U}\left\langle\frac{\delta\ell}{\delta
w},\eta\right\rangle d^{n}x\right|_{0}^{T}$ (6.2)

When the section $\phi\in\mathcal{P}$, the first integral is equal to zero and
that leaves us with

$\displaystyle 0=\left.\int_{\partial U}\left\langle\frac{\delta L}{\delta
w},\eta\right\rangle d^{n}x\right|_{0}^{T}.$ (6.3)

The importance of this is that it allows us to find integrals of motion for
the Euler–Poincaré system and one of them is the exterior derivative of the
one-form part of $\frac{\delta\ell}{\delta w}$. For fluids, it is known as
_vorticity_. In [32], it was shown how Noether’s theorem can be used to derive
local conservation of vorticity. The same calculation is repeated here and it
starts by equating the following time derivative to zero

$\displaystyle 0=\frac{d}{dt}\left\langle\frac{\delta\ell}{\delta
w},\eta\right\rangle=\frac{d}{dt}\int_{U}\frac{\delta\ell}{\delta w}\cdot
dx\wedge\left(\eta\makebox[11.99998pt][c]{\mbox{\rule{6.00006pt}{0.72229pt}\rule{0.72229pt}{6.00006pt}}}\rho
d^{n}x\right),$ (6.4)

then substituting the vector field $\eta$ with $\frac{1}{\rho}\left(\star
d\left(\boldsymbol{\Psi}^{\flat}\right)\right)^{\sharp}\cdot\nabla$, where
$\boldsymbol{\Psi}$ is an arbitrary function. The operator
$\star:\Lambda^{k}\to\Lambda^{N-k}$ is the Hodge star operator and it is a
linear operator, which maps the $k$ forms on $N$ dimensional manifold to $n-k$
forms provided that $0\leq k\leq n$. Meanwhile, the operators
$\flat:\Lambda^{k}\to\left(\Lambda^{k}\right)^{\ast}$ and
$\sharp:\left(\Lambda^{k}\right)^{\ast}\to\Lambda^{k}$ are musical
isomorphisms. In $2$ and $3$ dimensions, $\left(\star
d\left(\Psi^{\flat}\right)\right)^{\sharp}=\mathrm{curl}(\boldsymbol{\Psi})$.
The calculation carried out yields

$\displaystyle
d\left(\partial_{t}+\pounds_{\theta}\right)\left(\frac{1}{\rho}\frac{\delta\ell}{\delta
w}\cdot dx\right)$
$\displaystyle=\left(\partial_{t}+\pounds_{\theta}\right)d\left(\frac{1}{\rho}\frac{\delta\ell}{\delta
w}\cdot dx\right)=0,$ (6.5)

Replacing $d\left(\frac{1}{\rho}\frac{\delta\ell}{\delta w}\cdot dx\right)$ by
$\omega\cdot dS$, where the differential form $dS$ is defined, by components,
by $dS_{i}=\frac{1}{2}\epsilon_{i_{1},i_{2},\dots,i_{N}}dx^{i_{j}}\wedge
dx^{i_{k}}$ , where $\epsilon_{i_{1},i_{2},\dots,i_{N}}$ is the Levi-Civita
symbol. Then equation (6.5), becomes

$\displaystyle\left(\partial_{t}+\pounds_{\theta}\right)\left(\omega\cdot
dS\right)=\left(\partial_{t}\omega+\left(\omega\cdot\nabla\right)\theta-\left(\theta\cdot\nabla\right)\omega+\omega\nabla\cdot\theta\right)\cdot
dS=0.$ (6.6)

In two or three dimensions, it is written in terms of curl operator

$\displaystyle\left(\partial_{t}+\pounds_{\theta}\right)\left(\mathrm{curl}\left(\frac{1}{\rho}\frac{\delta\ell}{\delta
w}\right)\cdot dS\right)=0.$

Going back to equation (6.6), the term
$\left(\omega\cdot\nabla\right)\theta-\left(\theta\cdot\nabla\right)\omega$ is
the Lie derivative $\pounds_{\theta}\omega$, thus the coefficient of the two-
form $dS$, $\omega$, evolves according to the following directional derivative

$\displaystyle\frac{d}{dt}\omega=\left(\partial_{t}+\pounds_{\theta}\right)\omega=-\omega\nabla\cdot\theta.$

Using the fact that the coefficient of the continuity equation can be written
as

$\displaystyle\frac{d}{dt}\rho=-\rho\nabla\cdot\theta,$ (6.7)

and we use that in the equation
$\frac{d}{dt}\omega=\frac{1}{\rho}\frac{d}{dt}\rho$, which leads to

$\displaystyle\frac{d}{dt}\left(\frac{\omega}{\rho}\right)=0,$ (6.8)

and it means that the quantity $\frac{\omega}{\rho}$ is a conserved quantity.
In section 6.3, this quantity plays an important role for investigating the
nonlinear stability of the deep learning training problem.

A special version of Noether’s theorem for Euler–Poincaré equations and it is
known as _Kelvin-Noether’s theorem_ [56].It includes Kelvin’s circulation
theorem where it states that the circulation of a closed loop transported in
an barotropic fluid is a constant. Given a circulation map
$\kappa:\mathfrak{C}\to\mathfrak{X}(Q)^{\ast\ast}$, where $\mathfrak{C}$ is
the space of continuous loops in $Q$. The circulation map is defined by

$\displaystyle\left\langle\kappa(C),\alpha\right\rangle=\oint_{C}\frac{\alpha}{\rho},$
(6.9)

where $C\in\mathfrak{C}$ is a loop and $\alpha\in\mathfrak{X}(Q)^{\ast}$ is a
one-form density and dividing it by $\rho$ renders it a one-form. This path
integral is known as _circulation_ [56]. According to Kelvin-Noether’s theorem
(theorem $4.1$ in [56]), and given the quantity, known as _Kelvin-Noether
quantity_ , $I:\mathfrak{C}\times\mathfrak{X}(Q)\to\mathbb{R}$ defined by

$\displaystyle I(C,v,t)=\left\langle\kappa(C),\frac{\delta\ell}{\delta
w}\right\rangle,$

and its time evolution is governed by the following differential equation

$\displaystyle\frac{d}{dt}I(C,w,t)=\left\langle\kappa(C(t)),\rho\diamond\frac{\delta\ell}{\delta\rho}\right\rangle,$

where $C(t)=C_{0}g(t)^{-1}$ is a loop in $\mathfrak{C}$,
$C_{0}\in\mathfrak{C}$ and $g$ is the solution of $\dot{g}=R_{g^{-1}}w$ and
$w,\rho$ satisfy the Euler–Poincaré equation. Then, Kelvin’s circulation
theorem can be stated in terms of Kelvin-Noether’s theorem as

$\displaystyle I(C,w,t)=\oint_{C}\frac{1}{\rho}\left(\frac{\delta\ell}{\delta
w}\right).$

###### Lemma 6.1 (Conservation of Kelvin’s circulation for deep learning).

Given $(w,\rho)$, which satisfies the Euler–Poincaré equation (3.21), and let
the diffeomorphism $g$ is the solution of $\dot{g}=w\circ g(t)$. With the loop
$C(t)=g\circ C_{0}$, then Kelvin’s circulation

$\displaystyle I(C,w,t)=\oint_{C}\frac{1}{\rho}\left(\frac{\delta\ell}{\delta
w}\right),$ (6.10)

is conserved.

###### Proof.

In order to prove the conservation of $I$, Kelvin’s circulation is rewritten
as

$\displaystyle
I(C,w,t)=\oint_{C_{0}}(g)^{\ast}\left(\frac{1}{\rho}\left(\frac{\delta\ell}{\delta
w}\right)\right)=\oint_{C_{0}}\frac{1}{\rho_{0}}(g)^{\ast}\left(\frac{\delta\ell}{\delta
w}\right),$ (6.11)

Taking the time derivative of $I$ and noticing that
$\frac{d}{dt}g^{\ast}(\frac{\delta\ell}{\delta w})$ is the Lie derivative and
it constitute the left–hand side of the Euler–Poincaré equation (3.21), then

$\displaystyle\frac{d}{dt}I=\frac{d}{dt}\oint_{C_{0}}\frac{1}{\rho_{0}}(g)^{\ast}\left(\frac{\delta\ell}{\delta
w}\right)$
$\displaystyle=\frac{d}{dt}\oint_{C_{0}}\frac{1}{\rho_{0}}g^{\ast}\left(\rho\diamond\frac{\delta\ell}{\delta\rho}\right)$
(6.12)
$\displaystyle=\oint_{C}\frac{1}{\rho}\left(\rho\nabla\frac{\delta\ell}{\delta\rho}\cdot